Handle older FlashInfer FP8 MoE signatures

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

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

.buildkite/ci_config_intel.yaml

@@ -2,17 +2,17 @@ name: vllm_intel_ci
 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"
   - "requirements/common.txt"
   - "requirements/xpu.txt"
-  - "requirements/build/cuda.txt"
-  - "requirements/test/cuda.txt"
   - "setup.py"
   - "csrc/"
   - "cmake/"
 run_all_exclude_patterns:
-  - "docker/Dockerfile."
   - "csrc/cpu/"
   - "csrc/rocm/"
   - "cmake/hipify.py"

.buildkite/hardware_tests/amd.yaml

@@ -6,6 +6,7 @@ steps:
   # differ ci_base is rebuilt and pushed automatically.
   - label: "AMD: :docker: ensure ci_base"
     key: ensure-ci-base-amd
+    soft_fail: false
     depends_on: []
     device: amd_cpu
     no_plugin: true
@@ -26,6 +27,7 @@ steps:
 
   - label: "AMD: :docker: build test image and artifacts"
     key: image-build-amd
+    soft_fail: false
     depends_on:
       - ensure-ci-base-amd
     device: amd_cpu

.buildkite/hardware_tests/cpu.yaml

@@ -53,7 +53,7 @@ steps:
   - tests/models/language/pooling/
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
@@ -68,13 +68,15 @@ steps:
   - vllm/v1/sample/ops/topk_topp_triton.py
   - vllm/v1/sample/ops/topk_topp_sampler.py
   - tests/v1/sample/test_topk_topp_sampler.py
+  - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
       VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
       # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
-      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
+      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp
+      pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
 
 - label: CPU-Quantization Model Tests
   depends_on: []
@@ -89,11 +91,13 @@ steps:
   - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
   - tests/quantization/test_compressed_tensors.py
   - tests/quantization/test_cpu_wna16.py
+  - tests/quantization/test_cpu_w8a8.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
-      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
+      pytest -x -v -s tests/quantization/test_cpu_wna16.py
+      pytest -x -v -s tests/quantization/test_cpu_w8a8.py"
       
 - label: CPU-Distributed Tests (PP+TP)
   depends_on: []

.buildkite/intel_jobs/basic_correctness.yaml

@@ -23,4 +23,5 @@ steps:
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      pytest -v -s basic_correctness/test_cpu_offload.py &&
       pytest -v -s basic_correctness/test_mem.py::test_end_to_end'

.buildkite/intel_jobs/lora_intel.yaml

@@ -81,7 +81,9 @@ steps:
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
+      pytest -v -s lora/test_punica_ops.py::test_kernels &&
+      pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size &&
+      pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'
 
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45
@@ -128,10 +130,10 @@ steps:
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
       pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
       pytest -v -s lora/test_transformers_model.py &&
       pytest -v -s lora/test_chatglm3_tp.py &&
+      pytest -v -s lora/test_llama_tp.py::test_llama_lora &&
       pytest -s -v lora/test_minicpmv_tp.py'
 
 - label: LoRA Multimodal

.buildkite/intel_jobs/misc_intel.yaml

@@ -103,6 +103,31 @@ steps:
       pytest -v -s v1/kv_offload &&
       pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
 
+- label: NixlConnector PD accuracy (2 GPUs)
+  timeout_in_minutes: 60
+  num_devices: 2
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/v1/worker/kv_connector_model_runner_mixin.py
+    - tests/v1/kv_connector/nixl_integration/
+    - vllm/platforms/xpu.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
+
 - label: Regression
   key: regression
   timeout_in_minutes: 30
@@ -133,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/models_distributed_intel.yaml

@@ -0,0 +1,27 @@
+group: Models - Distributed
+depends_on: 
+  - image-build-xpu
+steps:
+- label: Distributed Model Tests (2 GPUs)
+  key: distributed-model-tests-2-gpus
+  timeout_in_minutes: 50
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+  - vllm/model_executor/model_loader/sharded_state_loader.py
+  - vllm/model_executor/models/
+  - tests/model_executor/model_loader/test_sharded_state_loader.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m "not slow_test"'

.buildkite/intel_jobs/models_multimodal_intel.yaml

@@ -22,7 +22,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av &&
       cd tests &&
       pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
       pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
@@ -47,8 +47,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
 
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
@@ -71,8 +70,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
       pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
 
@@ -96,7 +94,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av &&
       cd tests &&
       pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
 
@@ -121,11 +119,9 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av matplotlib ftfy &&
       pip install open-clip-torch --no-deps &&
       cd tests &&
       pytest -v -s models/multimodal/processing/test_tensor_schema.py
-      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
-      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
       --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
   parallelism: 4

.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
@@ -68,7 +93,6 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
         pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
@@ -120,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/ci-bake-rocm.sh

@@ -15,9 +15,10 @@ set -euo pipefail
 
 DEFAULT_REPO_SLUG="vllm-project/vllm"
 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
-DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base tools/install_torchcodec_rocm.sh tests/vllm_test_utils"
+DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
 DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
+DEFAULT_CI_BASE_METADATA_VERSION="1"
 IMAGE_EXISTED_BEFORE_BUILD=0
 
 TARGET=""
@@ -525,6 +526,22 @@ get_remote_image_label_with_retry() {
     return 0
 }
 
+remote_ci_base_metadata_is_current() {
+    local image_ref="$1"
+    local metadata_version=""
+
+    metadata_version=$(get_remote_image_label "${image_ref}" "vllm.ci_base.metadata_version")
+    [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]]
+}
+
+remote_ci_base_metadata_is_current_with_retry() {
+    local image_ref="$1"
+    local metadata_version=""
+
+    metadata_version=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.metadata_version")
+    [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]]
+}
+
 remote_image_exists() {
     local image_ref="$1"
     docker manifest inspect "${image_ref}" >/dev/null 2>&1
@@ -581,6 +598,7 @@ init_config() {
     CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
     CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
     CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
+    CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
     CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
     export PYTORCH_ROCM_ARCH
 
@@ -635,6 +653,10 @@ load_ci_hcl() {
     echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}"
 }
 
+init_bake_files() {
+    BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
+}
+
 compute_ci_base_hash_if_needed() {
     if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then
         return 0
@@ -676,12 +698,14 @@ configure_ci_base_image_refs() {
     fi
 
     content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}")
+    CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}"
     if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
         commit_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${BUILDKITE_COMMIT}")
-        CI_BASE_IMAGE_TAG_COMMIT="${commit_tag}"
-        export CI_BASE_IMAGE_TAG_COMMIT
     fi
+    CI_BASE_IMAGE_TAG_COMMIT_REF="${commit_tag}"
 
+    # *_REF is the logical tag recorded in metadata. *_EXTRA is only passed to
+    # bake when that tag is not already the primary tag, avoiding duplicates.
     if should_push_stable_ci_base_tag; then
         primary_tag="${content_tag}"
         CI_BASE_IMAGE_TAG_STABLE="${stable_tag}"
@@ -691,19 +715,33 @@ configure_ci_base_image_refs() {
     fi
     CI_BASE_IMAGE_TAG="${primary_tag}"
     if [[ "${primary_tag}" == "${content_tag}" ]]; then
-        CI_BASE_IMAGE_TAG_CONTENT=""
+        CI_BASE_IMAGE_TAG_CONTENT_EXTRA=""
     else
-        CI_BASE_IMAGE_TAG_CONTENT="${content_tag}"
+        CI_BASE_IMAGE_TAG_CONTENT_EXTRA="${content_tag}"
     fi
-    export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE
+    if [[ -n "${commit_tag}" && "${commit_tag}" != "${primary_tag}" ]]; then
+        CI_BASE_IMAGE_TAG_COMMIT_EXTRA="${commit_tag}"
+    else
+        CI_BASE_IMAGE_TAG_COMMIT_EXTRA=""
+    fi
+    export CI_BASE_IMAGE_TAG
+    export CI_BASE_IMAGE_TAG_COMMIT_EXTRA
+    export CI_BASE_IMAGE_TAG_CONTENT_EXTRA
+    export CI_BASE_IMAGE_TAG_CONTENT_REF
+    export CI_BASE_IMAGE_TAG_COMMIT_REF
+    export CI_BASE_IMAGE_TAG_STABLE
 
     if is_ci_base_target; then
         IMAGE_TAG="${primary_tag}"
         export IMAGE_TAG
 
         echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
-        if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then
-            echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}"
+        if [[ -n "${commit_tag}" ]]; then
+            if [[ "${commit_tag}" == "${primary_tag}" ]]; then
+                echo "ci_base commit image tag: ${commit_tag} (primary)"
+            else
+                echo "ci_base commit image tag: ${commit_tag}"
+            fi
         fi
         echo "ci_base content image tag: ${content_tag}"
         if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then
@@ -728,8 +766,8 @@ ci_base_candidate_refs() {
     printf '%s\n' \
         "${IMAGE_TAG:-}" \
         "${CI_BASE_IMAGE_TAG:-}" \
-        "${CI_BASE_IMAGE_TAG_COMMIT:-}" \
-        "${CI_BASE_IMAGE_TAG_CONTENT:-}" \
+        "${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}" \
+        "${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}" \
         "${CI_BASE_IMAGE_TAG_STABLE:-}" \
         | awk 'NF && !seen[$0]++'
 }
@@ -743,6 +781,10 @@ find_matching_ci_base_ref() {
         remote_image_exists "${candidate}" || continue
         candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash")
         if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
+            if ! remote_ci_base_metadata_is_current "${candidate}"; then
+                echo "Found matching ci_base content hash but stale metadata: ${candidate}" >&2
+                continue
+            fi
             printf '%s\n' "${candidate}"
             return 0
         fi
@@ -817,6 +859,10 @@ maybe_skip_existing_image() {
         if [[ -n "${remote_hash}" ]]; then
             echo "Remote ci_base content hash: ${remote_hash:0:16}..."
             if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
+                if ! remote_ci_base_metadata_is_current "${IMAGE_TAG}"; then
+                    echo "Content hashes match but ci_base metadata is stale; rebuilding to refresh metadata"
+                    return 0
+                fi
                 if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then
                     echo "ci_base tag refresh failed; rebuilding to push expected tags"
                     return 0
@@ -998,12 +1044,104 @@ prepare_git_cache_metadata() {
     fi
 }
 
+ci_base_metadata_pairs() {
+    local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
+    local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
+    local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
+    local content_files_hash=""
+    local base_image=""
+    local base_image_digest=""
+    local git_branch=""
+    local -a content_paths=()
+    local -a content_args=()
+
+    read -r -a content_paths <<< "${content_files}"
+    if [[ ${#content_paths[@]} -gt 0 ]]; then
+        content_files_hash=$(compute_content_hash "${content_paths[@]}")
+    fi
+    mapfile -t content_args < <(
+        get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}"
+    )
+
+    base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE")
+    if [[ -n "${base_image}" ]]; then
+        base_image_digest=$(resolve_image_digest "${base_image}")
+    fi
+    git_branch="${BUILDKITE_BRANCH:-${VLLM_BRANCH:-}}"
+
+    metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
+    metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}"
+    metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}"
+    metadata_pair "vllm.ci_base.content_files" "${content_files}"
+    metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")"
+    metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}"
+    metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}"
+    metadata_pair "vllm.ci_base.image.primary" "${CI_BASE_IMAGE_TAG:-}"
+    metadata_pair "vllm.ci_base.image.content" "${CI_BASE_IMAGE_TAG_CONTENT_REF:-${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}}"
+    metadata_pair "vllm.ci_base.image.commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}}"
+    metadata_pair "vllm.ci_base.image.stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
+    metadata_pair "vllm.ci_base.git_commit" "${BUILDKITE_COMMIT:-}"
+    metadata_pair "vllm.ci_base.git_branch" "${git_branch}"
+    metadata_pair "vllm.ci_base.vllm_branch" "${VLLM_BRANCH:-}"
+    metadata_pair "vllm.ci_base.stable_branch" "${CI_BASE_STABLE_BRANCH:-main}"
+
+    metadata_pair "vllm.rocm.base_image" "${base_image}"
+    metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}"
+    metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}"
+    metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
+    metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
+    metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
+    metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
+    metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
+    metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
+    metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
+    metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
+    metadata_pair "vllm.rocm.rocshmem_commit" "${ROCSHMEM_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_BRANCH")}"
+    metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")"
+    metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
+    metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
+    metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
+    metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
+    metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
+    metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
+
+    metadata_pair "vllm.buildkite.build_number" "${BUILDKITE_BUILD_NUMBER:-}"
+    metadata_pair "vllm.buildkite.build_id" "${BUILDKITE_BUILD_ID:-}"
+}
+
+write_ci_base_metadata_annotations() {
+    local metadata="$1"
+    local key=""
+    local value=""
+    local annotation=""
+
+    [[ -n "${metadata}" ]] || return 0
+    while IFS=$'\t' read -r key value; do
+        [[ -n "${key}" && -n "${value}" ]] || continue
+        annotation="manifest:${key}=${value}"
+        printf '    "%s",\n' "$(hcl_escape_string "${annotation}")"
+    done <<< "${metadata}"
+}
+
+write_ci_base_metadata_labels() {
+    local metadata="$1"
+    local key=""
+    local value=""
+
+    [[ -n "${metadata}" ]] || return 0
+    while IFS=$'\t' read -r key value; do
+        [[ -n "${key}" && -n "${value}" ]] || continue
+        printf '    "%s" = "%s"\n' \
+            "$(hcl_escape_string "${key}")" \
+            "$(hcl_escape_string "${value}")"
+    done <<< "${metadata}"
+}
+
 write_ci_base_label_override() {
     local target_name=""
+    local metadata=""
     local -a ci_base_targets=()
 
-    BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
-
     if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
         return 0
     fi
@@ -1019,16 +1157,23 @@ write_ci_base_label_override() {
         return 0
     fi
 
+    metadata=$(ci_base_metadata_pairs)
+
     : > "${CI_BASE_LABEL_OVERRIDE_PATH}"
     for target_name in "${ci_base_targets[@]}"; do
         cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
 target "${target_name}" {
   annotations = [
     "manifest:org.opencontainers.image.revision=",
+EOF
+        write_ci_base_metadata_annotations "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
+        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
   ]
   labels = {
     "org.opencontainers.image.revision" = ""
-    "vllm.ci_base.content_hash" = "${CI_BASE_CONTENT_HASH}"
+EOF
+        write_ci_base_metadata_labels "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
+        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
   }
 }
 
@@ -1036,7 +1181,7 @@ EOF
     done
 
     BAKE_FILES+=(-f "${CI_BASE_LABEL_OVERRIDE_PATH}")
-    echo "Appended ci_base content-hash label override for targets: ${ci_base_targets[*]}"
+    echo "Appended ci_base metadata label override for targets: ${ci_base_targets[*]}"
 }
 
 uses_rocm_csrc_cache() {
@@ -1119,6 +1264,18 @@ hcl_escape_string() {
     printf '%s' "${value}"
 }
 
+join_words() {
+    local IFS=" "
+    printf '%s' "$*"
+}
+
+metadata_pair() {
+    local key="$1"
+    local value="${2:-}"
+
+    printf '%s\t%s\n' "${key}" "${value}"
+}
+
 write_hcl_string_list() {
     local indent="$1"
     shift
@@ -1541,7 +1698,13 @@ confirm_remote_image_push() {
 
         remote_hash=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.content_hash")
         if [[ -n "${remote_hash}" && "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
-            return 0
+            if remote_ci_base_metadata_is_current_with_retry "${image_ref}"; then
+                return 0
+            fi
+
+            echo "Remote image exists with the expected ci_base content hash but stale metadata."
+            echo "  expected metadata version: ${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
+            return 1
         fi
 
         echo "Remote image exists but does not have the expected ci_base content hash."
@@ -1724,15 +1887,16 @@ main() {
     print_header
     validate_inputs
     load_ci_hcl
+    init_bake_files
     compute_ci_base_hash_if_needed
     configure_ci_base_image_refs
     maybe_skip_existing_image
     setup_builder
     prepare_git_cache_metadata
-    write_ci_base_label_override
     extract_dependency_pins
     write_rocm_build_arg_override
     compute_dependency_cache_keys
+    write_ci_base_label_override
     compute_rocm_csrc_content_hash_if_needed
     write_rocm_cache_override
     resolve_ci_base_dependency_targets

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

@@ -367,6 +367,20 @@ remove_docker_container() {
 }
 trap remove_docker_container EXIT
 
+# python_only_compile.sh runs `python setup.py develop` and needs the full repo tree
+# under /vllm-workspace (Dockerfile.rocm test stage: mkdir src && mv vllm).
+# The ROCm wheel artifact tarball only ships a thin tree (tests, etc.), so
+# artifact images cannot satisfy that test — use the full rocm/vllm-ci image.
+_cmd_probe="${VLLM_TEST_COMMANDS:-}"
+if [[ -z "${_cmd_probe}" ]]; then
+  _cmd_probe="$*"
+fi
+if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" == "1" && "${_cmd_probe}" == *python_only_compile.sh* ]]; then
+  echo "INFO: disabling VLLM_CI_USE_ARTIFACTS for python_only_compile (requires full /vllm-workspace tree)"
+  export VLLM_CI_USE_ARTIFACTS=0
+fi
+unset -v _cmd_probe
+
 if ! prepare_artifact_image; then
   echo "Using full ROCm CI image: ${image_name}"
   docker pull "${image_name}" || exit 1
@@ -426,6 +440,26 @@ fi
 
 echo "Final commands: $commands"
 
+# The ROCm test image often ships /vllm-workspace without .git (artifact tarball unpack).
+# tests/standalone_tests/python_only_compile.sh uses merge-base(HEAD, origin/main) for
+# wheels.vllm.ai; compute on the agent (full git checkout) and pass into the container.
+vllm_standalone_merge_base=""
+checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
+if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
+  checkout="."
+fi
+# Pass safe.directory per-command (-c) because buildkite runs will always fail
+# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
+if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+  vllm_standalone_merge_base="$(
+    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+  )"
+fi
+if [[ -z "${vllm_standalone_merge_base}" ]]; then
+  vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
+fi
+echo "INFO: passing VLLM_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
+
 MYPYTHONPATH="/vllm-workspace"
 
 container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
@@ -502,6 +536,20 @@ else
   echo "--- Single-node job"
   echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
 
+  ulimit_core_hard=$(ulimit -H -c)
+  if [[ "$ulimit_core_hard" == "unlimited" ]]; then
+    # docker run can't pass "unlimited" to --ulimit
+    ulimit_core_hard="-1"
+  fi
+   # Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled
+  coredump_flags="--ulimit core=0:$ulimit_core_hard"
+  if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then
+    # Works around https://github.com/rocm/rocm-systems/issues/6206
+    coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"'
+  else
+    echo "ROCm debug agent not enabled, coredumps are disabled in the test container."
+  fi
+
   docker run \
     --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
     $RDMA_FLAGS \
@@ -509,6 +557,7 @@ else
     --shm-size=16gb \
     --group-add "$render_gid" \
     --rm \
+    $coredump_flags \
     -e HF_TOKEN \
     -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
     -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
@@ -525,6 +574,7 @@ else
     -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
     -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
     -e "PYTORCH_ROCM_ARCH=" \
+    -e "VLLM_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}" \
     --name "${container_name}" \
     "${image_name}" \
     /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"

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

@@ -8,7 +8,7 @@ set -ex
 CORE_RANGE=${CORE_RANGE:-0-31}
 OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-31}
 
-export CMAKE_BUILD_PARALLEL_LEVEL=16
+export CMAKE_BUILD_PARALLEL_LEVEL=32
 
 # Setup cleanup
 remove_docker_container() {
@@ -37,7 +37,7 @@ function cpu_tests() {
     pytest -x -v -s tests/kernels/test_onednn.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
     pytest -x -v -s tests/kernels/core/test_cpu_activation.py
-    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
+    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
     pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
 
   # skip tests requiring model downloads if HF_TOKEN is not set

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

@@ -7,10 +7,49 @@ set -euox pipefail
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-48-95}
 NUMA_NODE=${NUMA_NODE:-1}
-IMAGE_NAME="cpu-test-$NUMA_NODE"
+AGENT_SLOT=${AGENT_SLOT:-}
+IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}"
 TIMEOUT_VAL=$1
 TEST_COMMAND=$2
 
+# Disk hygiene knobs. Reclaim space only once the Docker root filesystem crosses
+# DISK_USAGE_THRESHOLD percent, and cap the shared BuildKit cache at
+# BUILDKIT_CACHE_MAX so subsequent builds keep reusing the hottest layers.
+DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-70}
+BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-80GB}
+
+# Reclaim disk only when the host is under pressure. We trim (not purge) the
+# shared BuildKit cache so cross-job/cross-agent reuse stays intact, and only
+# touch dangling images; other agents' uniquely tagged images are left alone.
+prune_if_disk_pressure() {
+    local docker_root disk_usage
+    docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true)
+    if [ -z "$docker_root" ]; then
+        return 0
+    fi
+    disk_usage=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%')
+    if [ "${disk_usage:-0}" -gt "$DISK_USAGE_THRESHOLD" ]; then
+        echo "--- :broom: Disk usage ${disk_usage}% exceeds ${DISK_USAGE_THRESHOLD}%, reclaiming space"
+        docker image prune -f || true
+        docker builder prune -f --keep-storage="$BUILDKIT_CACHE_MAX" || true
+    else
+        echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune"
+    fi
+}
+
+# Always drop this agent's image once the job ends (the default builder never
+# uses it as a cache source, so removing it costs no rebuild speed), then
+# reclaim space if needed. Guard every docker call with `|| true` so the trap
+# never overrides the test's exit code.
+cleanup() {
+    docker image rm -f "$IMAGE_NAME" || true
+    prune_if_disk_pressure
+}
+trap cleanup EXIT
+
+# Free space up front so a nearly-full host doesn't fail the build.
+prune_if_disk_pressure
+
 # building the docker image
 echo "--- :docker: Building Docker image"
 docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .

.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 \
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
     -e CMDS \
     --name "${container_name}" \
     "${IMAGE}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
     >/dev/null
 } 9>/tmp/docker-pull.lock
 

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

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

.buildkite/scripts/run-rust-frontend-cargo-ci.sh

@@ -90,6 +90,16 @@ install_cargo_sort() {
   cargo binstall --no-confirm cargo-sort
 }
 
+install_cargo_deny() {
+  if command -v cargo-deny >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing cargo-deny"
+  install_cargo_binstall
+  cargo binstall --no-confirm cargo-deny
+}
+
 install_cargo_nextest() {
   if command -v cargo-nextest >/dev/null 2>&1; then
     return
@@ -142,6 +152,7 @@ PY
 
 run_style_clippy() {
   install_cargo_sort
+  install_cargo_deny
 
   log_section "Checking Rust formatting"
   cargo fmt --manifest-path rust/Cargo.toml --all -- --check
@@ -149,6 +160,13 @@ run_style_clippy() {
   log_section "Checking Cargo.toml ordering"
   cargo sort --workspace --check rust
 
+  log_section "Checking Rust dependency bans"
+  cargo deny \
+    --manifest-path rust/Cargo.toml \
+    check \
+    --config rust/deny.toml \
+    bans
+
   log_section "Running clippy"
   cargo clippy \
     --manifest-path rust/Cargo.toml \

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh

@@ -33,6 +33,14 @@ if [[ -n "${ATTENTION_BACKEND:-}" ]]; then
   EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
 fi
 
+# ROCm: run eager to avoid intermittent HIP-graph decode corruption.
+# See https://github.com/ROCm/clr/issues/279
+# TODO(aarushjain29): Revert after TheRock 7.14
+if command -v rocm-smi &> /dev/null || command -v amd-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
+  echo "ROCm platform detected: adding --enforce-eager to avoid HIP-graph decode corruption"
+  EXTRA_ARGS+=(--enforce-eager)
+fi
+
 cleanup() {
   if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
     kill "${SERVER_PID}" 2>/dev/null || true

.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_areas/basic_correctness.yaml

@@ -16,3 +16,9 @@ steps:
   - pytest -v -s basic_correctness/test_mem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      depends_on:
+      - image-build-amd

.buildkite/test_areas/benchmarks.yaml

@@ -11,6 +11,11 @@ steps:
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
 
 - label: Attention Benchmarks Smoke Test (B200)
   key: attention-benchmarks-smoke-test-b200

.buildkite/test_areas/cuda.yaml

@@ -2,8 +2,8 @@ group: CUDA
 depends_on: 
   - image-build
 steps:
-- label: Platform Tests (CUDA)
-  key: platform-tests-cuda
+- label: Platform Tests
+  key: platform-tests
   timeout_in_minutes: 15
   device: h200_18gb
   source_file_dependencies:

.buildkite/test_areas/disaggregated.yaml

@@ -13,6 +13,20 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      device: mi300_4
+      timeout_in_minutes: 110
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - tests/v1/kv_connector/nixl_integration/
+        - vllm/platforms/rocm.py
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
   key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
   timeout_in_minutes: 30
@@ -36,6 +50,19 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      device: mi300_4
+      timeout_in_minutes: 50
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - tests/v1/kv_connector/nixl_integration/
+        - vllm/platforms/rocm.py
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
   key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
@@ -48,6 +75,19 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      device: mi300_4
+      timeout_in_minutes: 110
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - tests/v1/kv_connector/nixl_integration/
+        - vllm/platforms/rocm.py
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
   key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
@@ -60,6 +100,19 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      device: mi300_4
+      timeout_in_minutes: 60
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - tests/v1/kv_connector/nixl_integration/
+        - vllm/platforms/rocm.py
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
   key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
@@ -103,6 +156,20 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+  mirror:
+    amd:
+      device: mi300_2
+      timeout_in_minutes: 60
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - vllm/v1/worker/kv_connector_model_runner_mixin.py
+        - tests/v1/kv_connector/nixl_integration/
+        - vllm/platforms/rocm.py
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
   key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus

.buildkite/test_areas/distributed.yaml

@@ -37,6 +37,21 @@ steps:
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
   - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
+  mirror:
+    amd:
+      device: mi300_2
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/
+      - vllm/engine/
+      - vllm/executor/
+      - vllm/worker/worker_base.py
+      - vllm/v1/engine/
+      - vllm/v1/worker/
+      - tests/v1/distributed
+      - tests/entrypoints/openai/test_multi_api_servers.py
+      - vllm/platforms/rocm.py
 
 - label: Distributed Compile + RPC Tests (2 GPUs)
   key: distributed-compile-rpc-tests-2-gpus
@@ -159,8 +174,8 @@ steps:
   # test multi-node TP with multiproc executor (simulated on single node)
   - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
 
-- label: Distributed Tests (8 GPUs)(H100)
-  key: distributed-tests-8-gpus-h100
+- label: Distributed Tests (8xH100)
+  key: distributed-tests-8xh100
   timeout_in_minutes: 10
   device: h100
   num_devices: 8
@@ -180,8 +195,8 @@ steps:
   # test with torchrun tp=2 and dp=4 with ep
   - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
 
-- label: Distributed Tests (4 GPUs)(A100)
-  key: distributed-tests-4-gpus-a100
+- label: Distributed Tests (4xA100)
+  key: distributed-tests-4xa100
   device: a100
   optional: true
   num_devices: 4
@@ -195,8 +210,8 @@ steps:
   - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - pytest -v -s -x lora/test_mixtral.py
 
-- label: Distributed Tests (2 GPUs)(H100)
-  key: distributed-tests-2-gpus-h100
+- label: Distributed Tests (2xH100-2xMI300)
+  key: distributed-tests-2xh100-2xmi300
   timeout_in_minutes: 15
   device: h100
   optional: true
@@ -210,15 +225,15 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
     - pytest -v -s tests/distributed/test_packed_tensor.py
 
-- label: Distributed Tests (2 GPUs)(B200)
-  key: distributed-tests-2-gpus-b200
+- label: Distributed Tests (2xB200)
+  key: distributed-tests-2xb200
   device: b200-k8s
   optional: true
   working_dir: "/vllm-workspace/"
   num_devices: 2
   commands:
     - pytest -v -s tests/distributed/test_context_parallel.py
-    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
+    - pytest -v -s tests/distributed/test_nccl_symm_mem.py
     - pytest -v -s tests/v1/distributed/test_dbo.py
     - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
 

.buildkite/test_areas/e2e_integration.yaml

@@ -2,8 +2,8 @@ group: E2E Integration
 depends_on: 
   - image-build
 steps:
-- label: DeepSeek V2-Lite Sync EPLB Accuracy
-  key: deepseek-v2-lite-sync-eplb-accuracy
+- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100)
+  key: deepseek-v2-lite-sync-eplb-accuracy-4xh100
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -12,8 +12,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100)
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -22,8 +22,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200)
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200
   timeout_in_minutes: 60
   device: b200-k8s
   optional: true

.buildkite/test_areas/engine.yaml

@@ -112,6 +112,11 @@ steps:
   commands:
     # Only run tests that need exactly 2 GPUs
     - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
+  mirror:
+    amd:
+      device: mi300_2
+      depends_on:
+      - image-build-amd
 
 - label: V1 e2e (4 GPUs)
   key: v1-e2e-4-gpus

.buildkite/test_areas/entrypoints.yaml

@@ -29,6 +29,8 @@ steps:
   mirror:
     amd:
       device: mi325_1
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -40,10 +42,12 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/entrypoints/serve
+  - tests/entrypoints/scale_out
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
+  - pytest -v -s entrypoints/scale_out
   mirror:
     amd:
       device: mi325_1

.buildkite/test_areas/expert_parallelism.yaml

@@ -14,6 +14,16 @@ steps:
   commands:
   - pytest -v -s distributed/test_eplb_algo.py
   - pytest -v -s distributed/test_eplb_utils.py
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/eplb
+      - tests/distributed/test_eplb_algo.py
+      - tests/distributed/test_eplb_utils.py
+      - vllm/platforms/rocm.py
 
 - label: EPLB Execution # 17min
   key: eplb-execution

.buildkite/test_areas/kernels.yaml

@@ -47,8 +47,10 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py
 
 - label: Deepseek V4 Kernel Test (B200)
   key: deepseek-v4-kernel-test-b200
@@ -272,8 +274,8 @@ steps:
     - pytest -v -s kernels/helion/
 
  
-- label: Kernels FP8 MoE Test (1 H100)
-  key: kernels-fp8-moe-test-1-h100
+- label: Kernels FP8 MoE Test (1xH100)
+  key: kernels-fp8-moe-test-1xh100
   timeout_in_minutes: 90
   device: h100
   num_devices: 1
@@ -289,8 +291,8 @@ steps:
     - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
     - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
 
-- label: Kernels FP8 MoE Test (2 H100s)
-  key: kernels-fp8-moe-test-2-h100s
+- label: Kernels FP8 MoE Test (2xH100)
+  key: kernels-fp8-moe-test-2xh100
   timeout_in_minutes: 90
   device: h100
   num_devices: 2

.buildkite/test_areas/lm_eval.yaml

@@ -28,7 +28,8 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
 
-# - label: LM Eval Large Models (4 GPUs)(A100)
+# - label: LM Eval Large Models (4xA100)
+#   key: lm-eval-large-models-4xa100
 #   device: a100
 #   optional: true
 #   num_devices: 4
@@ -40,8 +41,8 @@ steps:
 #   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
 
-- label: LM Eval Large Models (4 GPUs)(H100)
-  key: lm-eval-large-models-4-gpus-h100
+- label: LM Eval Large Models (4xH100)
+  key: lm-eval-large-models-4xh100
   device: h100
   optional: true
   num_devices: 4
@@ -53,8 +54,8 @@ steps:
     - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
     - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
 
-- label: LM Eval Small Models (B200)
-  key: lm-eval-small-models-b200
+- label: LM Eval Small Models (1xB200)
+  key: lm-eval-small-models-1xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -64,8 +65,21 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
 
-- label: LM Eval Large Models (B200, EP)
-  key: lm-eval-large-models-b200-ep
+- label: LM Eval Small Models Distributed (2xB200)
+  key: lm-eval-small-models-distributed-2xb200
+  timeout_in_minutes: 120
+  device: b200-k8s
+  num_devices: 2
+  optional: true
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
+
+- label: LM Eval Large Models EP (2xB200)
+  key: lm-eval-large-models-ep-2xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -76,8 +90,8 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt
 
-- label: LM Eval Qwen3.5 Models (B200)
-  key: lm-eval-qwen3-5-models-b200
+- label: LM Eval Qwen3.5 Models (2xB200)
+  key: lm-eval-qwen3-5-models-2xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -93,8 +107,8 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
 
-- label: LM Eval Large Models (H200)
-  key: lm-eval-large-models-h200
+- label: LM Eval Large Models (8xH200)
+  key: lm-eval-large-models-8xh200
   timeout_in_minutes: 60
   device: h200
   optional: true
@@ -192,8 +206,8 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
 
-- label: GPQA Eval (GPT-OSS) (H100)
-  key: gpqa-eval-gpt-oss-h100
+- label: GPQA Eval (GPT-OSS) (2xH100)
+  key: gpqa-eval-gpt-oss-2xh100
   timeout_in_minutes: 120
   device: h100
   optional: true
@@ -206,8 +220,8 @@ steps:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
 
-- label: GPQA Eval (GPT-OSS) (B200)
-  key: gpqa-eval-gpt-oss-b200
+- label: GPQA Eval (GPT-OSS) (2xB200)
+  key: gpqa-eval-gpt-oss-2xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -226,6 +240,8 @@ steps:
   device: dgx-spark
   optional: true
   num_devices: 1
+  depends_on:
+  - arm64-image-build
   source_file_dependencies:
   - csrc/
   - vllm/model_executor/layers/quantization

.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
 
@@ -224,6 +224,16 @@ steps:
     - 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
     # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
     - 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
+  mirror:
+    amd:
+      device: mi325_1
+      source_file_dependencies:
+      - vllm/entrypoints
+      - vllm/multimodal
+      - examples/
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 - label: Metrics, Tracing (2 GPUs)
   key: metrics-tracing-2-gpus
@@ -250,6 +260,12 @@ steps:
       'opentelemetry-exporter-otlp>=1.26.0' \
       'opentelemetry-semantic-conventions-ai>=0.4.1'"
   - pytest -v -s v1/tracing
+  mirror:
+    amd:
+      device: mi325_2
+      depends_on:
+      - image-build-amd
+      optional: true
 
 - label: Python-only Installation
   key: python-only-installation
@@ -262,6 +278,16 @@ steps:
   - setup.py
   commands:
   - bash standalone_tests/python_only_compile.sh
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 20
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/standalone_tests/python_only_compile.sh
+      - setup.py
+      - vllm/platforms/rocm.py
 
 - label: Async Engine, Inputs, Utils, Worker
   device: h200_35gb
@@ -344,7 +370,7 @@ steps:
   - pytest -v -s test_ray_env.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
+  - pytest -v -s reasoning
   - pytest -v -s tool_parsers
   - pytest -v -s tokenizers_
   - pytest -v -s parser

.buildkite/test_areas/model_executor.yaml

@@ -23,3 +23,16 @@ steps:
     # calls that the signal method cannot interrupt.
     - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
     - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/engine/arg_utils.py
+      - vllm/config/model.py
+      - vllm/model_executor
+      - tests/model_executor
+      - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py

.buildkite/test_areas/model_runner_v2.yaml

@@ -18,9 +18,7 @@ steps:
   - set -x
   - export VLLM_USE_V2_MODEL_RUNNER=1
   - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  # This requires eager until we sort out CG correctness issues.
-  # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
-  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
   - pytest -v -s v1/e2e/general/test_context_length.py
   - pytest -v -s v1/e2e/general/test_min_tokens.py
   # Temporary hack filter to exclude ngram spec decoding based tests.

.buildkite/test_areas/models_basic.yaml

@@ -6,7 +6,6 @@ steps:
   key: basic-models-tests-initialization
   timeout_in_minutes: 45
   device: h200_18gb
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/test_initialization.py
@@ -14,8 +13,6 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
-  mirror:
-    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   device: h200_35gb
@@ -31,8 +28,6 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
-  mirror:
-    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   device: h200_35gb
@@ -45,6 +40,11 @@ steps:
   - tests/models/test_registry.py
   commands:
     - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
 
 - label: Basic Models Test (Other CPU) # 5min
   key: basic-models-test-other-cpu

.buildkite/test_areas/models_language.yaml

@@ -14,7 +14,10 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
   mirror:
-    torch_nightly: {}
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
 
 - label: Language Models Tests (Extra Standard) %N
   key: language-models-tests-extra-standard
@@ -31,7 +34,21 @@ steps:
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
   mirror:
-    torch_nightly: {}
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/layers/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - tests/models/language/pooling/test_embedding.py
+      - tests/models/language/generation/test_common.py
+      - tests/models/language/pooling/test_classification.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: Language Models Tests (Hybrid) %N
   key: language-models-tests-hybrid
@@ -48,7 +65,6 @@ steps:
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
   mirror:
-    torch_nightly: {}
     amd:
       device: mi325_1
       timeout_in_minutes: 90

.buildkite/test_areas/models_multimodal.yaml

@@ -10,7 +10,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2"
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
@@ -27,7 +26,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
   mirror:
@@ -44,7 +42,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
@@ -61,7 +58,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
     - pytest models/multimodal/generation/test_memory_leak.py -m core_model
@@ -83,7 +79,6 @@ steps:
   - tests/models/registry.py
   device: cpu-medium
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
 
 - label: Multi-Modal Processor # 44min
@@ -95,7 +90,6 @@ steps:
   - tests/models/multimodal
   - tests/models/registry.py
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
@@ -109,6 +103,17 @@ steps:
   - vllm/v1/core/
   commands:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/multimodal/
+      - vllm/inputs/
+      - vllm/v1/core/
+      - vllm/platforms/rocm.py
+      - vllm/model_executor/model_loader/
 
 - label: Multi-Modal Models (Extended Generation 1)
   key: multi-modal-models-extended-generation-1
@@ -118,7 +123,6 @@ steps:
   - tests/models/multimodal/generation
   - tests/models/multimodal/test_mapping.py
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
@@ -135,7 +139,6 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
 
 - label: Multi-Modal Models (Extended Generation 3)
@@ -146,7 +149,6 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
 
 - label: Multi-Modal Models (Extended Pooling)

.buildkite/test_areas/pytorch.yaml

@@ -195,3 +195,11 @@ steps:
   - requirements/test/nightly-torch.txt
   commands:
   - bash standalone_tests/pytorch_nightly_dependency.sh
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - requirements/test/nightly-torch.txt
+      - vllm/platforms/rocm.py

.buildkite/test_areas/rust_frontend.yaml

@@ -26,7 +26,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
   # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
 
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
@@ -46,7 +46,7 @@ steps:
   - vllm/v1/engine/
   - tests/utils.py
   # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
-  - tests/entrypoints/serve/disagg/test_serving_tokens.py
+  - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
   - tests/entrypoints/serve/instrumentator/test_basic.py
   - tests/entrypoints/serve/instrumentator/test_metrics.py
   # - tests/entrypoints/serve/dev/test_sleep.py
@@ -55,7 +55,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
   - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
-  - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
+  - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
   - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
   # - pytest -v -s entrypoints/serve/dev/test_sleep.py
 

.buildkite/test_areas/spec_decode.yaml

@@ -94,6 +94,8 @@ steps:
     amd:
       device: mi325_1
       timeout_in_minutes: 65
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/weight_loading.yaml

@@ -13,6 +13,13 @@ steps:
   - tests/weight_loading
   commands:
     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
+  mirror:
+    amd:
+      device: mi300_2
+      depends_on:
+      - image-build-amd
+      commands:
+      - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
 
 # - label: Weight Loading Multiple GPU - Large Models # optional
 #   working_dir: "/vllm-workspace/tests"

.github/CODEOWNERS

@@ -3,7 +3,7 @@
 
 # This lists cover the "core" components of vLLM that require careful review
 /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
-/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
+/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
@@ -11,7 +11,7 @@
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/model_loader @22quinn
-/vllm/model_executor/layers/batch_invariant.py @yewentao256 
+/vllm/model_executor/layers/batch_invariant.py @yewentao256
 /vllm/ir @ProExpertProg
 /vllm/kernels/ @ProExpertProg @tjtanaa
 /vllm/kernels/helion @ProExpertProg @zou3519
@@ -23,7 +23,7 @@
 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
 /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
-/vllm/config/cache.py @heheda12345
+/vllm/config/cache.py @heheda12345 @ivanium
 
 # Config utils
 /vllm/config/utils.py @hmellor
@@ -67,16 +67,17 @@
 /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety @vadiklyutiy
 /vllm/v1/attention/backends/triton_attn.py @tdoublep
 /vllm/v1/attention/backends/gdn_attn.py @ZJY0516 @vadiklyutiy
-/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
+/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
 /vllm/v1/sample @22quinn @houseroad @njhill
 /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
 /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
-/vllm/v1/kv_cache_interface.py @heheda12345
+/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium
 /vllm/v1/kv_offload @ApostaC @orozery
+/vllm/v1/simple_kv_offload @ivanium
 /vllm/v1/engine @njhill
 /vllm/v1/executor @njhill
 /vllm/v1/worker @njhill
-/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
+/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium
 
 # Model runner V2
 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256
@@ -103,13 +104,14 @@
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
-/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
+/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
 /tests/weight_loading @mgoin @youkaichao @yewentao256
 /tests/lora @jeejeelee
 /tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91
 /tests/v1/kv_connector/nixl_integration @NickLucche
-/tests/v1/kv_connector @ApostaC @orozery
+/tests/v1/kv_connector @ApostaC @orozery @ivanium
 /tests/v1/kv_offload @ApostaC @orozery
+/tests/v1/simple_kv_offload @ivanium
 /tests/v1/determinism @yewentao256
 /tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
 /tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning

.github/actionlint.yaml

@@ -3,3 +3,5 @@
 self-hosted-runner:
   labels:
     - vllm-runners
+    # Not yet in actionlint's known-label set.
+    - macos-26

.github/workflows/issue_autolabel.yml

@@ -327,7 +327,7 @@ jobs:
                 message: 'CC {users} for ROCm-related issue',
               },
               mistral: {
-                users: ['patrickvonplaten', 'juliendenize', 'andylolu2'],
+                users: ['patrickvonplaten', 'juliendenize', 'andylolu2', 'NickLucche'],
                 message: 'CC {users} for Mistral-related issue',
               },
               // Add more label -> user mappings here

.github/workflows/macos-smoke-test.yml

@@ -11,13 +11,25 @@ permissions:
 
 jobs:
   macos-m1-smoke-test:
-    runs-on: macos-latest
+    # macos-26 (the supported target) is still a preview runner, so gate on GA
+    # macos-15 and keep macos-26 non-blocking.
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: macos-15
+            required: true
+          - os: macos-26
+            required: false
+    name: macos-m1-smoke-test (${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    continue-on-error: ${{ !matrix.required }}
     timeout-minutes: 30
 
     steps:
-      - uses: actions/checkout@v6.0.1
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
 
-      - uses: astral-sh/setup-uv@v7
+      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
         with:
           enable-cache: true
           cache-dependency-glob: |
@@ -72,14 +84,11 @@ jobs:
           # Test health endpoint
           curl -f http://localhost:8000/health
 
-          # Test completion
-          curl -f http://localhost:8000/v1/completions \
+          # Long prompt: hits the split-KV path that short prompts skip (#46769).
+          PAYLOAD=$(python -c "import json; print(json.dumps({'model': 'Qwen/Qwen3-0.6B', 'prompt': 'The quick brown fox jumps over the lazy dog. ' * 24, 'max_tokens': 16}))")
+          curl -f --max-time 120 http://localhost:8000/v1/completions \
             -H "Content-Type: application/json" \
-            -d '{
-              "model": "Qwen/Qwen3-0.6B",
-              "prompt": "Hello",
-              "max_tokens": 5
-            }'
+            -d "$PAYLOAD"
 
           # Cleanup
           kill "$SERVER_PID"

.github/workflows/pre-commit.yml

@@ -48,7 +48,7 @@ jobs:
     if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
     runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
-    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
     - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
       with:
         python-version: "3.12"

.pre-commit-config.yaml

@@ -131,6 +131,19 @@ repos:
         --python-version, "3.12",
       ]
       files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$
+    - id: pip-compile
+      alias: pip-compile-cpu
+      name: pip-compile-cpu
+      args: [
+        requirements/test/cuda.in,
+        -o, requirements/test/cpu.txt,
+        --index-strategy, unsafe-best-match,
+        --torch-backend, cpu,
+        --python-platform, x86_64-manylinux_2_28,
+        --python-version, "3.12",
+      ]
+      files: ^requirements/(common|cpu|test/(cuda|cpu))\.(in|txt)$
+      exclude: ^requirements/test/cuda\.txt$
     - id: pip-compile
       alias: pip-compile-docs
       name: pip-compile-docs

CMakeLists.txt

@@ -140,6 +140,21 @@ if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
     WITH_SOABI)
 endif()
 
+#
+# fs_io extension (pure CXX; must stay above the non-CUDA device branch
+# so CPU builds define the target before the early return).
+# GIL-releasing filesystem helpers for FileSystemTierManager.
+#
+if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
+  define_extension_target(
+    fs_io_C
+    DESTINATION vllm
+    LANGUAGE CXX
+    SOURCES csrc/fs_io.cpp
+    USE_SABI 3.11
+    WITH_SOABI)
+endif()
+
 #
 # Forward the non-CUDA device extensions to external CMake scripts.
 #
@@ -270,6 +285,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   #
   set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result -Wno-unused-value")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result -Wno-unused-value")
+
+  # When using LTO then *.cpp files must be compiled with same compiler as used linker
+  # So if HIP uses clang linker we also must use it
+  # Otherwise symbols will be missing from .so
+  if (CMAKE_CXX_FLAGS MATCHES "\-flto")
+    if(NOT CMAKE_CXX_COMPILER_ID STREQUAL CMAKE_HIP_COMPILER_ID)
+      message(FATAL_ERROR "LTO is enabled for ROCm build, but the C++ compiler (${CMAKE_CXX_COMPILER_ID}) and HIP compiler (${CMAKE_HIP_COMPILER_ID}) are different which is not supported. "
+                          "Please ensure they are same by setting CXX=${CMAKE_HIP_COMPILER} environment variable. Or alternatively disable LTO.")
+    endif()
+   endif()
 endif()
 
 #
@@ -325,9 +350,7 @@ endif()
 if(VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_EXT_SRC
     "csrc/torch_bindings.cpp"
-    "csrc/custom_quickreduce.cu"
-    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu")
+    "csrc/custom_quickreduce.cu")
 
   message(STATUS "Enabling C extension.")
   define_extension_target(
@@ -355,6 +378,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
+    "csrc/libtorch_stable/cuda_view.cu"
+    "csrc/libtorch_stable/cuda_utils_kernels.cu"
     "csrc/libtorch_stable/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
@@ -374,14 +399,30 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels_fused.cu"
     "csrc/libtorch_stable/custom_all_reduce.cu"
     "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
 
+  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
+     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
+     CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
+
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+
+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
+
+    endif()
+  endif()
+
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -416,8 +457,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     FetchContent_MakeAvailable(cutlass)
 
     list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/libtorch_stable/cuda_view.cu"
-      "csrc/libtorch_stable/cuda_utils_kernels.cu"
       "csrc/libtorch_stable/cutlass_extensions/common.cpp"
       "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
@@ -498,6 +537,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${VLLM_STABLE_EXT_SRC}"
       CUDA_ARCHS "${CUDA_ARCHS}")
 
+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_STABLE_EXT_SRC
+        "csrc/libtorch_stable/cooperative_topk.cu")
+      set_gencode_flags_for_srcs(
+        SRCS "csrc/libtorch_stable/cooperative_topk.cu"
+        CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
+    endif()
+
     # Only build Marlin kernels if we are building for at least some compatible archs.
     # Keep building Marlin for 9.0 as there are some group sizes and shapes that
     # are not supported by Machete yet.
@@ -843,9 +890,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   endif()
 
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
     set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
@@ -1040,25 +1087,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     USE_SABI 3
     WITH_SOABI)
 
+  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
+  # This ensures we only use C-shim APIs available in PyTorch 2.11.
+  # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+  # which is currently set to 2.11.
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
   # Needed to use cuda/hip APIs from C-shim
   if(VLLM_GPU_LANG STREQUAL "CUDA")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.11.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.11.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020B000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    if(COOPERATIVE_TOPK_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_COOPERATIVE_TOPK=1)
+    endif()
     # Needed by CUTLASS kernels
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
   elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.10.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.10.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020A000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
   endif()
 
@@ -1266,25 +1312,20 @@ define_extension_target(
   USE_SABI 3
   WITH_SOABI)
 
+# Set TORCH_TARGET_VERSION for stable ABI compatibility.
+# This ensures we only use C-shim APIs available in PyTorch 2.11.
+# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+# which is currently set to 2.11.
+target_compile_definitions(_moe_C_stable_libtorch PRIVATE
+  TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
 # Needed to use cuda/hip APIs from C-shim
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.11.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.11.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020B000000000000ULL)
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
   # Needed by CUTLASS kernels
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE
     CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 elseif(VLLM_GPU_LANG STREQUAL "HIP")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.10.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.10.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020A000000000000ULL)
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
 endif()
 

benchmarks/benchmark_topk_topp.py

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

benchmarks/kernels/benchmark_fused_collective.py

@@ -80,13 +80,17 @@ _FI_MAX_SIZES = {
     2: 64 * MiB,  # 64MB
     4: 64 * MiB,  # 64MB
     8: 64 * MiB,  # 64MB
+    16: 64 * MiB,  # 64MB (multi-node)
 }
 
 # Global workspace tensors for FlashInfer (keyed by backend name)
 _FI_WORKSPACES: dict = {}
 
-# Backends to benchmark
-FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
+# Backends to benchmark. trtllm is single-node only and can hang cross-node, so
+# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl.
+FLASHINFER_BACKENDS = [
+    b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b
+]
 
 
 def setup_flashinfer_workspace(
@@ -995,7 +999,10 @@ def main():
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
 
-    device = torch.device(f"cuda:{rank}")
+    # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per
+    # node) map to a valid local GPU; falls back to global rank single-node.
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    device = torch.device(f"cuda:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
 

benchmarks/kernels/benchmark_moe.py

@@ -391,16 +391,19 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int
         config = dict(zip(keys, config_values))
         configs.append(config)
 
-    # Remove configs that are not compatible with fp8 block quantization
-    # BLOCK_SIZE_K must be a multiple of block_k
-    # BLOCK_SIZE_N must be a multiple of block_n
+    # Drop configs incompatible with fp8 block quantization. A tile must align
+    # to the quant-block scale grid, i.e. tile and block must divide one
+    # another. The kernel indexes scales per element (offs_bn // group_n,
+    # k_start // group_k), so a tile narrower than the block (e.g. N=64 with
+    # block_n=128) is valid -- and often faster at small batch. An exact
+    # multiple was required before, which dropped those smaller tiles entirely.
     if block_quant_shape is not None and not use_fp16:
         block_n, block_k = block_quant_shape[0], block_quant_shape[1]
         for config in configs[:]:
-            if (
-                config["BLOCK_SIZE_K"] % block_k != 0
-                or config["BLOCK_SIZE_N"] % block_n != 0
-            ):
+            bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"]
+            n_aligned = bn % block_n == 0 or block_n % bn == 0
+            k_aligned = bk % block_k == 0 or block_k % bk == 0
+            if not (n_aligned and k_aligned):
                 configs.remove(config)
     return configs
 

benchmarks/kernels/benchmark_moe_defaults.py

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

benchmarks/kernels/benchmark_paged_attention.py

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

benchmarks/kernels/benchmark_selective_state_update.py

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

benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py

@@ -7,6 +7,7 @@ import time
 import numpy as np
 import torch
 
+from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import set_random_seed
 
@@ -14,17 +15,15 @@ from vllm.utils.torch_utils import set_random_seed
 try:
     from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
 except (ImportError, AttributeError) as e:
-    print("ERROR: CPU fused MoE operations are not available on this platform.")
-    print("This benchmark requires x86 CPU with proper vLLM CPU extensions compiled.")
-    print(
-        "The cpu_fused_moe kernel is typically available on Linux x86_64 "
-        "with AVX2/AVX512."
-    )
     print(f"Import error: {e}")
     sys.exit(1)
 
 # ISA selection following test_cpu_fused_moe.py pattern
-ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
+ISA_CHOICES = ["vec"]
+if torch.cpu._is_amx_tile_supported():
+    ISA_CHOICES.append("amx")
+if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+    ISA_CHOICES.append("neon")
 
 
 @torch.inference_mode()
@@ -145,7 +144,7 @@ if __name__ == "__main__":
         "--isa",
         type=str,
         choices=ISA_CHOICES,
-        default=ISA_CHOICES[0],
+        default="vec",
         help=f"ISA to use (available: {ISA_CHOICES})",
     )
     parser.add_argument("--seed", type=int, default=0)

cmake/cpu_extension.cmake

@@ -24,7 +24,10 @@ set (ENABLE_NUMA TRUE)
 # Check the compile flags
 #
 if(MACOSX_FOUND)
+    # Apple clang needs -Xpreprocessor to enable OpenMP. No runtime link is
+    # needed: _C is a dynamic_lookup bundle and resolves libomp from torch.
     list(APPEND CXX_COMPILE_FLAGS
+        "-Xpreprocessor" "-fopenmp"
         "-DVLLM_CPU_EXTENSION")
 else()
     list(APPEND CXX_COMPILE_FLAGS
@@ -166,12 +169,13 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
-    if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
+    if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0)
         message(FATAL_ERROR
-            "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
+            "VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'")
     endif()
     # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
-    # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
+    # by default; set -DVLLM_RVV_VLEN=0 to force scalar RISC-V build.
+    # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
     if(NOT DEFINED VLLM_RVV_VLEN)
         # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
         if(EXISTS /proc/cpuinfo)
@@ -325,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")
@@ -423,6 +427,8 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
         "csrc/cpu/activation_lut_bf16.cpp"
+        "csrc/cpu/cpu_tanhf_neon.hpp"
+        "csrc/cpu/cpu_fused_moe.cpp"
         ${VLLM_EXT_SRC})
 endif()
 

cmake/external_projects/deepgemm.cmake

@@ -29,7 +29,8 @@ if(DEEPGEMM_SRC_DIR)
 else()
   # Keep in sync with tools/install_deepgemm.sh
   set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
-  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+  # NOTE: This is currently targeting nv-dev branch due to sm120 support
+  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

cmake/external_projects/fmha_sm100.cmake

@@ -17,7 +17,7 @@ else()
   FetchContent_Declare(
     fmha_sm100
     GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
@@ -36,13 +36,38 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
 
 install(FILES
   "${FMHA_SM100_PY_ROOT}/__init__.py"
+  "${FMHA_SM100_PY_ROOT}/api.py"
+  "${FMHA_SM100_PY_ROOT}/bench_utils.py"
+  "${FMHA_SM100_PY_ROOT}/jit.py"
   "${FMHA_SM100_PY_ROOT}/sparse.py"
+  "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
   DESTINATION vllm/third_party/fmha_sm100
   COMPONENT fmha_sm100)
 
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
+  DESTINATION vllm/third_party/fmha_sm100/csrc
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
   DESTINATION vllm/third_party/fmha_sm100/cute
   COMPONENT fmha_sm100
   PATTERN "__pycache__" EXCLUDE
   PATTERN "*.pyc" EXCLUDE
   PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
+          GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/cpu/activation.cpp

@@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
       });
 }
 
+void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
+  int num_tokens = input.numel() / input.size(-1);
+  int d = input.size(-1);
+
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
+    CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
+    activation_kernel<scalar_t, gelu_tanh_act, false>(
+        num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
+    CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
+  });
+}
+
 void gelu_new(torch::Tensor& out, torch::Tensor& input) {
   int num_tokens = input.numel() / input.size(-1);
   int d = input.size(-1);

csrc/cpu/cpu_attn_impl.hpp

@@ -124,7 +124,7 @@ struct AttentionMetadata {
         workitem_group_num(workitem_group_num),
         reduction_item_num(reduction_item_num),
         reduction_split_num(reduction_split_num),
-        thread_num(omp_get_max_threads()),
+        thread_num(cpu_utils::get_max_threads()),
         effective_thread_num(thread_num),
         split_kv_q_token_num_threshold(split_kv_q_token_num_threshold),
         attention_scratchpad_size_per_thread(0),
@@ -405,7 +405,7 @@ class AttentionScheduler {
   torch::Tensor schedule(const ScheduleInput& input) const {
     const bool causal = input.causal;
     const bool is_dynamic_causal = input.dynamic_causal != nullptr;
-    const int32_t thread_num = omp_get_max_threads();
+    const int32_t thread_num = cpu_utils::get_max_threads();
     const int64_t cache_size = cpu_utils::get_available_l2_size();
     const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
     const int32_t kv_len_alignment = input.kv_block_alignment;
@@ -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;
     }
@@ -887,12 +929,10 @@ struct VecTypeTrait<c10::BFloat16> {
   using vec_t = vec_op::BF16Vec16;
 };
 
-#if !defined(__powerpc__)
 template <>
 struct VecTypeTrait<c10::Half> {
   using vec_t = vec_op::FP16Vec16;
 };
-#endif
 
 template <typename T>
 void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
@@ -1425,7 +1465,7 @@ class AttentionMainLoop {
 
  public:
   void operator()(const AttentionInput* input) {
-    const int thread_num = omp_get_max_threads();
+    const int thread_num = cpu_utils::get_max_threads();
     TORCH_CHECK_EQ(input->metadata->thread_num, thread_num);
     std::atomic<int32_t> guard_counter(0);
     std::atomic<int32_t>* guard_counter_ptr = &guard_counter;

csrc/cpu/cpu_attn_vsx.hpp

@@ -50,7 +50,16 @@ FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
   b1 = (__vector float)vec_mergel(zeros, raw);
 }
 
-// Note: c10::Half (FP16) is not supported on PowerPC architecture
+// [3] Half (FP16) Specialization
+template <>
+FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
+                                                __vector float& b0,
+                                                __vector float& b1) {
+  vec_op::FP16Vec8 fp16_vec(p);
+  vec_op::FP32Vec8 fp32_vec(fp16_vec);
+  b0 = fp32_vec.reg.val[0];
+  b1 = fp32_vec.reg.val[1];
+}
 
 template <int32_t M, typename kv_cache_t>
 FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(

csrc/cpu/cpu_fused_moe.cpp

@@ -1,5 +1,3 @@
-#include <sleef.h>
-
 #include "cpu/cpu_types.hpp"
 #include "cpu/utils.hpp"
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -16,6 +14,18 @@
   #define AMX_DISPATCH(...) case cpu_utils::ISA::AMX:
 #endif
 
+#if defined(ARM_BF16_SUPPORT)
+  #include "cpu/micro_gemm/cpu_micro_gemm_neon.hpp"
+  #define NEON_DISPATCH(...)                                         \
+    case cpu_utils::ISA::NEON: {                                     \
+      using gemm_t =                                                 \
+          cpu_micro_gemm::MicroGemm<cpu_utils::ISA::NEON, scalar_t>; \
+      return __VA_ARGS__();                                          \
+    }
+#else
+  #define NEON_DISPATCH(...) case cpu_utils::ISA::NEON:
+#endif
+
 #define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...)                          \
   [&] {                                                               \
     switch (ISA_TYPE) {                                               \
@@ -25,6 +35,7 @@
             cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \
         return __VA_ARGS__();                                         \
       }                                                               \
+        NEON_DISPATCH(__VA_ARGS__)                                    \
       default: {                                                      \
         TORCH_CHECK(false, "Invalid CPU ISA type.");                  \
       }                                                               \
@@ -59,10 +70,12 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                        const int32_t input_stride,
                        const int32_t output_stride) {
   using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
+#if !defined(__aarch64__)
   // For GPT-OSS interleaved gate-up weights
   alignas(64) static int32_t index[16] = {0,  2,  4,  6,  8,  10, 12, 14,
                                           16, 18, 20, 22, 24, 26, 28, 30};
   vec_op::INT32Vec16 index_vec(index);
+#endif
   vec_op::FP32Vec16 gate_up_max_vec(7.0);
   vec_op::FP32Vec16 up_min_vec(-7.0);
   vec_op::FP32Vec16 alpha_vec(1.702);
@@ -72,8 +85,15 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
 
   for (int32_t m = 0; m < m_size; ++m) {
     for (int32_t n = 0; n < n_size; n += 32) {
+      // Note: AdvSIMD does not support gather loads
+#if defined(__aarch64__)
+      vec_op::FP32Vec16 gate_vec(vec_op::uninit);
+      vec_op::FP32Vec16 up_vec(vec_op::uninit);
+      vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
+#else
       vec_op::FP32Vec16 gate_vec(input + n, index_vec);
       vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
+#endif
       gate_vec = gate_vec.min(gate_up_max_vec);
       up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
       auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
@@ -174,7 +194,7 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
       auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
       // Note: can't use fast_exp form because diffusiongemma will generate
       // wrong results
-      vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
+      auto tanh_vec = inner_vec.tanh();
       auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
       auto gated_output_fp32 = up_vec * gelu_tanh;
       scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
@@ -240,13 +260,14 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
   constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
   constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
   constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
+  constexpr bool pack_a = gemm_t::PackA;
   static_assert(gemm_n_tile_size % 16 == 0);
 
   TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
   TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
   TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
 
-  const int32_t thread_num = omp_get_max_threads();
+  const int32_t thread_num = cpu_utils::get_max_threads();
 
   const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
       gemm_m_tile_size * input_size_13 * sizeof(scalar_t));
@@ -266,12 +287,18 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
 
   const int32_t w2_input_tile_size = cpu_utils::round_up<64>(
       gemm_m_tile_size * input_size_2 * sizeof(scalar_t));
+  // use w2 input buffer only when we need to pack input
+  const int32_t w2_input_buffer_size =
+      pack_a ? cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 *
+                                       sizeof(scalar_t))
+             : 0;
 
   const int32_t w2_n_tile_size = [&]() {
     const int64_t cache_size = cpu_utils::get_available_l2_size();
-    // input tile + weight
+    // input tile + optional packed input + weight
     const int32_t n_size_cache_limit =
-        (cache_size - w2_input_tile_size) / (input_size_2 * sizeof(scalar_t));
+        (cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) /
+        (input_size_2 * sizeof(scalar_t));
     const int32_t n_size_thread_limit =
         output_size_2 / std::max(1, thread_num / topk_num);
     const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
@@ -324,6 +351,9 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
   const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;
   w13_thread_buffer_offset += w13_output_buffer_size;
 
+  const int32_t w2_input_buffer_offset = w13_thread_buffer_offset;
+  w13_thread_buffer_offset += w2_input_buffer_size;
+
   // Weighted sum thread buffer
   const int32_t ws_output_buffer_size =
       cpu_utils::round_up<64>(output_size_2 * sizeof(float));
@@ -403,7 +433,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
       gemm_t gemm;
 
       const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t);
-      const int32_t w13_n_group_stride = 16 * input_size_13;
+      const int32_t w13_n_group_stride =
+          gemm_t::WeightOCGroupSize * input_size_13;
       const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
 
       for (;;) {
@@ -466,8 +497,23 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
              token_idx += gemm_m_tile_size) {
           const int32_t actual_token_num =
               std::min(gemm_m_tile_size, curr_token_num - token_idx);
-          // copy inputs
-          {
+
+          scalar_t* __restrict__ curr_w13_gemm_input_buffer = nullptr;
+          if constexpr (pack_a) {
+            // copy and pack inputs
+            curr_w13_gemm_input_buffer = w13_input_buffer;
+            const scalar_t* w13_input_rows[gemm_m_tile_size];
+            for (int32_t i = 0; i < actual_token_num; ++i) {
+              w13_input_rows[i] =
+                  input + curr_expand_token_id_buffer[i] * input_size_13;
+            }
+            gemm_t::pack_input_from_rows(w13_input_rows,
+                                         curr_w13_gemm_input_buffer,
+                                         actual_token_num, input_size_13);
+            curr_expand_token_id_buffer += actual_token_num;
+          } else {
+            // copy inputs
+            curr_w13_gemm_input_buffer = curr_w13_input_buffer;
             scalar_t* __restrict__ curr_w13_input_buffer_iter =
                 curr_w13_input_buffer;
             for (int32_t i = 0; i < actual_token_num; ++i) {
@@ -499,14 +545,12 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
             scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1;
             scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0;
             scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1;
-            scalar_t* __restrict__ curr_w13_input_buffer_iter =
-                curr_w13_input_buffer;
             float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer;
             float* __restrict__ w13_output_buffer_1_iter =
                 w13_output_buffer + actual_n_tile_size / 2;
             for (int32_t i = 0; i < actual_n_tile_size;
                  i += min_w13_n_tile_size) {
-              gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_0_iter,
+              gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_0_iter,
                         w13_output_buffer_0_iter, actual_token_num,
                         input_size_13, input_size_13, w13_n_group_stride,
                         actual_n_tile_size, false);
@@ -519,7 +563,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
                 w13_bias_ptr_0_iter += gemm_n_tile_size;
               }
 
-              gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_1_iter,
+              gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_1_iter,
                         w13_output_buffer_1_iter, actual_token_num,
                         input_size_13, input_size_13, w13_n_group_stride,
                         actual_n_tile_size, false);
@@ -572,7 +616,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
       gemm_t gemm;
 
       const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
-      const int32_t w2_n_group_stride = 16 * input_size_2;
+      const int32_t w2_n_group_stride =
+          gemm_t::WeightOCGroupSize * input_size_2;
 
       for (;;) {
         int32_t task_id = counter_ptr->acquire_counter();
@@ -611,13 +656,30 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
              token_idx += gemm_m_tile_size) {
           const int32_t actual_token_num =
               std::min(gemm_m_tile_size, curr_token_num - token_idx);
+          scalar_t* __restrict__ curr_w2_gemm_input_buffer =
+              curr_w13_gemm_output_buffer;
+          if constexpr (pack_a) {
+            uint8_t* __restrict__ thread_buffer =
+                thread_buffer_start + thread_id * w13_thread_buffer_offset;
+            scalar_t* __restrict__ w2_input_buffer =
+                reinterpret_cast<scalar_t*>(thread_buffer +
+                                            w2_input_buffer_offset);
+            curr_w2_gemm_input_buffer = w2_input_buffer;
+            const scalar_t* w2_input_rows[gemm_m_tile_size];
+            for (int32_t i = 0; i < actual_token_num; ++i) {
+              w2_input_rows[i] = curr_w13_gemm_output_buffer + i * input_size_2;
+            }
+            gemm_t::pack_input_from_rows(w2_input_rows,
+                                         curr_w2_gemm_input_buffer,
+                                         actual_token_num, input_size_2);
+          }
 
           scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr;
           scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr;
           float* __restrict__ curr_w2_gemm_output_buffer_iter =
               curr_w2_gemm_output_buffer;
           for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
-            gemm.gemm(curr_w13_gemm_output_buffer, w2_weight_ptr_iter,
+            gemm.gemm(curr_w2_gemm_input_buffer, w2_weight_ptr_iter,
                       curr_w2_gemm_output_buffer_iter, actual_token_num,
                       input_size_2, input_size_2, w2_n_group_stride,
                       output_size_2, false);

csrc/cpu/cpu_tanhf_neon.hpp

@@ -0,0 +1,128 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#ifndef CPU_TANHF_NEON_HPP
+#define CPU_TANHF_NEON_HPP
+
+#include <cstdint>
+#include <arm_neon.h>
+
+namespace vec_op {
+
+namespace {
+
+struct TanhfConstants {
+  float32x4_t special_bound;
+  float32x4_t two;
+  float32x4_t c0;
+  float32x4_t c2;
+  int32x4_t exponent_bias;
+  float c1;
+  float c3;
+  float two_over_ln2;
+  float c4;
+  float ln2_hi;
+  float ln2_lo;
+};
+
+const TanhfConstants kTanhfConstants = {
+    // 9.01, above which tanhf rounds to 1 (or -1 for  negative).
+    .special_bound = vdupq_n_f32(0x1.205966p+3f),
+    .two = vdupq_n_f32(0x1.0p+1f),
+    .c0 = vdupq_n_f32(0x1.fffffep-2f),
+    .c2 = vdupq_n_f32(0x1.555736p-5f),
+    .exponent_bias = vdupq_n_s32(0x3f800000),
+    .c1 = 0x1.5554aep-3f,
+    .c3 = 0x1.12287cp-7f,
+    .two_over_ln2 = 0x1.715476p+1f,
+    .c4 = 0x1.6b55a2p-10f,
+    .ln2_hi = 0x1.62e4p-1f,
+    .ln2_lo = 0x1.7f7d1cp-20f,
+};
+
+// Return the ptr but hide it's value from the compiler so accesses
+// through it can't be optimised based on contents.
+template <typename T>
+inline const T* ptr_barrier(const T* ptr) {
+  const T* opaque_ptr = ptr;
+  __asm__("" : "+r"(opaque_ptr));
+  return opaque_ptr;
+}
+
+// Check whether any lanes in the mask are set
+inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
+
+// e^2x - 1 inline helper
+inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
+  float32x2_t ln2 = vld1_f32(&d->ln2_hi);
+  float32x4_t lane_consts = vld1q_f32(&d->c1);
+
+  // Reduce argument: f in [-ln2/2, ln2/2], i is exact.
+  float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
+  int32x4_t i = vcvtq_s32_f32(j);
+  float32x4_t f = vaddq_f32(x, x);
+  f = vfmsq_lane_f32(f, j, ln2, 0);
+  f = vfmsq_lane_f32(f, j, ln2, 1);
+
+  // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
+  float32x4_t f2 = vmulq_f32(f, f);
+  float32x4_t f4 = vmulq_f32(f2, f2);
+  float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
+  float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
+  float32x4_t poly = vfmaq_f32(p01, f2, p23);
+  poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
+  poly = vfmaq_f32(f, f2, poly);
+
+  // scale = 2^i
+  int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
+  float32x4_t scale = vreinterpretq_f32_s32(u);
+  return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
+}
+
+// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
+inline float32x4_t special_case(float32x4_t x, float32x4_t q,
+                                uint32x4_t special) {
+  const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
+
+  float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
+  uint32x4_t ix = vreinterpretq_u32_f32(x);
+  uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
+  uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
+  uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
+  float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
+  return vbslq_f32(special, special_y, y);
+}
+
+}  // namespace
+
+// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
+// AdvSIMD)
+// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
+//
+// Approximation for single-precision vector tanh(x), using a simplified
+// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
+// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
+inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
+  const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
+
+  // tanh(x) = (e^2x - 1) / (e^2x + 1)
+  // q = e^2x -1
+  float32x4_t q = e2xm1f_inline(x, d);
+
+  // Check for special cases
+  uint32x4_t special = vcagtq_f32(x, d->special_bound);
+
+  // Fall back to vectorised special case for any lanes which would cause
+  // expm1 to overflow
+  if (any_u32(special)) {
+    return special_case(x, q, special);
+  }
+
+  // Complete fast path if no special lanes
+  // tanh(x) = q / (q+2)
+  return vdivq_f32(q, vaddq_f32(q, d->two));
+}
+
+}  // namespace vec_op
+
+#endif  // CPU_TANHF_NEON_HPP

csrc/cpu/cpu_types.hpp

@@ -25,4 +25,20 @@
   #include <omp.h>
 #endif
 
+#include <c10/util/Exception.h>
+
+namespace cpu_utils {
+// Without OpenMP the omp pragmas compile to serial loops, so report 1: kernels
+// that barrier on the thread count would otherwise deadlock.
+inline int get_max_threads() {
+#ifdef _OPENMP
+  return omp_get_max_threads();
+#else
+  TORCH_WARN_ONCE(
+      "vLLM CPU was built without OpenMP; running single-threaded.");
+  return 1;
+#endif
+}
+}  // namespace cpu_utils
+
 #endif

csrc/cpu/cpu_types_arm.hpp

@@ -3,6 +3,8 @@
 
 #include <arm_neon.h>
 
+#include "cpu/cpu_tanhf_neon.hpp"
+
 #include <torch/all.h>
 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
   explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
 
   explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
+
+  FORCE_INLINE FP32Vec4 tanh() const {
+    return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
+  }
 };
 
 struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
     reg.val[1] = Vectorized<float>(data.val[1]);
   }
 
+  FORCE_INLINE FP32Vec8 tanh() const {
+    FP32Vec8 r(uninit);
+    r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
+    r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
+    return r;
+  }
+
   FORCE_INLINE float reduce_sum() const noexcept {
     float answer = 0;
     std::plus<VectorizedT> add;
@@ -497,6 +510,35 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
     reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
   };
 
+  FORCE_INLINE FP32Vec16 tanh() const {
+    FP32Vec16 r(uninit);
+    r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
+    r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
+    r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
+    r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
+    return r;
+  }
+
+  static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
+                                         FP32Vec16& odd) noexcept {
+    const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
+    const float32x4x2_t x23 =
+        vuzpq_f32(vld1q_f32(ptr + 8), vld1q_f32(ptr + 12));
+    const float32x4x2_t x45 =
+        vuzpq_f32(vld1q_f32(ptr + 16), vld1q_f32(ptr + 20));
+    const float32x4x2_t x67 =
+        vuzpq_f32(vld1q_f32(ptr + 24), vld1q_f32(ptr + 28));
+
+    even.reg.val[0] = VectorizedT(x01.val[0]);
+    even.reg.val[1] = VectorizedT(x23.val[0]);
+    even.reg.val[2] = VectorizedT(x45.val[0]);
+    even.reg.val[3] = VectorizedT(x67.val[0]);
+    odd.reg.val[0] = VectorizedT(x01.val[1]);
+    odd.reg.val[1] = VectorizedT(x23.val[1]);
+    odd.reg.val[2] = VectorizedT(x45.val[1]);
+    odd.reg.val[3] = VectorizedT(x67.val[1]);
+  }
+
   FORCE_INLINE FP32Vec16 operator+(const FP32Vec16& b) const noexcept {
     FP32Vec16 r(uninit);
     r.reg.val[0] = reg.val[0] + b.reg.val[0];
@@ -515,6 +557,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
     return r;
   }
 
+  FORCE_INLINE FP32Vec16 operator-() const noexcept {
+    FP32Vec16 r(uninit);
+    r.reg.val[0] = reg.val[0].neg();
+    r.reg.val[1] = reg.val[1].neg();
+    r.reg.val[2] = reg.val[2].neg();
+    r.reg.val[3] = reg.val[3].neg();
+    return r;
+  }
+
   FORCE_INLINE FP32Vec16 operator*(const FP32Vec16& b) const noexcept {
     FP32Vec16 r(uninit);
     r.reg.val[0] = reg.val[0] * b.reg.val[0];
@@ -933,4 +984,4 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
 
 inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); };
 
-};  // namespace vec_op
+};  // namespace vec_op

csrc/cpu/cpu_types_riscv_defs.hpp

@@ -3,13 +3,17 @@
 
 // VLEN-to-LMUL mapping for RISC-V Vector extension.
 //
-// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
-//   VLEN=128: LMUL_128=m1,  LMUL_256=m2,  LMUL_512=m4,  LMUL_1024=m8
-//   VLEN=256: LMUL_128=mf2, LMUL_256=m1,  LMUL_512=m2,  LMUL_1024=m4
+// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data.
+// LMUL_64 is used by 8-lane int8/uint8 vectors.
+//   VLEN=128:
+//     LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
+//   VLEN=256:
+//     LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
 
 #include <riscv_vector.h>
 
 #if __riscv_v_min_vlen == 128
+  #define LMUL_64 mf2
   #define LMUL_128 m1
   #define LMUL_256 m2
   #define LMUL_512 m4
@@ -17,6 +21,7 @@
   #define BOOL_256 b16
   #define BOOL_512 b8
 #elif __riscv_v_min_vlen == 256
+  #define LMUL_64 mf4
   #define LMUL_128 mf2
   #define LMUL_256 m1
   #define LMUL_512 m2
@@ -41,6 +46,16 @@
 
 // ---- Semantic fixed-vector typedefs (named by element count) ----
 
+// uint8 / int8
+typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t
+    __attribute__((riscv_rvv_vector_bits(64)));
+typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t
+    __attribute__((riscv_rvv_vector_bits(64)));
+
+// int16
+typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+
 // float16
 typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
     __attribute__((riscv_rvv_vector_bits(128)));

csrc/cpu/cpu_types_scalar.hpp

@@ -363,6 +363,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     return FP32Vec16(ret);
   }
 
+  FP32Vec16 tanh() const {
+    f32x16_t ret;
+    unroll_loop<int, VEC_ELEM_NUM>(
+        [&ret, this](int i) { ret.val[i] = std::tanh(reg.val[i]); });
+    return FP32Vec16(ret);
+  }
+
   float reduce_sum() const {
     float result = 0.0f;
     unroll_loop<int, VEC_ELEM_NUM>(

csrc/cpu/cpu_types_vsx.hpp

@@ -13,10 +13,10 @@ namespace vec_op {
 struct fp8_e4m3_tag {};
 struct fp8_e5m2_tag {};
 
-// FIXME: FP16 is not fully supported in Torch-CPU
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
 
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -34,6 +34,87 @@ struct fp8_e5m2_tag {};
 #define FORCE_INLINE __attribute__((always_inline)) inline
 
 namespace {
+
+FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
+  const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
+  const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
+  const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
+  const __vector unsigned int bias_adj = {112, 112, 112, 112};
+  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
+  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
+
+  __vector unsigned int s = (x & mask_sign) << 16;
+  __vector unsigned int e = (x & mask_exp) >> 10;
+  __vector unsigned int m = (x & mask_mant) << 13;
+
+  __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
+
+  __vector unsigned int e_normal = e + bias_adj;
+  e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
+
+  return (__vector float)(s | (e << 23) | m);
+}
+
+FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
+  __vector unsigned int in = (__vector unsigned int)f_in;
+
+  const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
+                                              0x80000000, 0x80000000};
+  const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
+                                             0x7F800000};
+  const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
+                                              0x007FFFFF, 0x007FFFFF};
+
+  const __vector signed int bias_adj = {112, 112, 112, 112};
+  const __vector signed int zero = {0, 0, 0, 0};
+  const __vector signed int max_exp = {31, 31, 31, 31};
+  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
+  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
+
+  __vector unsigned int s = (in & mask_sign_32) >> 16;
+  __vector unsigned int e_u = (in & mask_exp_32) >> 23;
+
+  __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
+
+  __vector signed int e_s = (__vector signed int)e_u;
+  e_s = vec_sub(e_s, bias_adj);
+  e_s = vec_max(e_s, zero);
+  e_s = vec_min(e_s, max_exp);
+  __vector unsigned int e_normal = (__vector unsigned int)e_s;
+
+  __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
+
+  const __vector unsigned int one_v = {1, 1, 1, 1};
+  const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
+
+  __vector unsigned int round_bit = (in >> 12) & one_v;
+  __vector unsigned int sticky = in & mask_sticky;
+  __vector unsigned int m = (in & mask_mant_32) >> 13;
+  __vector unsigned int lsb = m & one_v;
+
+  // Round up if: round_bit && (sticky || lsb)
+  __vector __bool int sticky_nonzero =
+      vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
+  __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
+  __vector __bool int round_up =
+      vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
+
+  m = vec_sel(m, m + one_v, round_up);
+
+  const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
+  const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
+  __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
+  __vector __bool int would_overflow_to_inf =
+      vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
+  __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
+  e_final = vec_sel(e_final, e_inc, mant_overflows);
+  m = vec_and(m, mant_mask);
+  e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
+  m = vec_sel(m, mant_mask, would_overflow_to_inf);
+
+  return s | (e_final << 10) | m;
+}
+
 template <typename T, T... indexes, typename F>
 constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
   (f(std::integral_constant<T, indexes>{}), ...);
@@ -89,6 +170,19 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
   }
 };
 
+struct FP16Vec8 : public Vec<FP16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+
+  __vector signed short reg;
+
+  explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
+  explicit FP16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    *reinterpret_cast<__vector signed short*>(ptr) = reg;
+  }
+};
+
 struct FP16Vec16 : public Vec<FP16Vec16> {
   constexpr static int VEC_ELEM_NUM = 16;
   ss16x8x2_t reg;
@@ -124,13 +218,11 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
   ss16x8x2_t reg;
 
   explicit BF16Vec16(const void* ptr) {
-    // Load 256 bits in two parts
     reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
     reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
   }
 
   explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
-
   explicit BF16Vec16(const FP32Vec16&);
 
   void save(void* ptr) const {
@@ -142,20 +234,16 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
   void save(void* ptr, const int elem_num) const {
     const int clamped_elem = std::max(0, std::min(elem_num, 16));
 
-    // Calculate elements to store in each 128-bit part (8 elements each)
     const int elements_val0 = std::min(clamped_elem, 8);
     const int elements_val1 = std::max(clamped_elem - 8, 0);
 
-    // Convert elements to bytes (2 bytes per element)
     const size_t bytes_val0 = elements_val0 * sizeof(signed short);
     const size_t bytes_val1 = elements_val1 * sizeof(signed short);
 
     signed short* dest = static_cast<signed short*>(ptr);
-    // Store the first part using vec_xst_len
     if (bytes_val0 > 0) {
       vec_xst_len(reg.val[0], dest, bytes_val0);
     }
-    // Store the second part if needed
     if (bytes_val1 > 0) {
       vec_xst_len(reg.val[1], dest + elements_val0, bytes_val1);
     }
@@ -238,6 +326,15 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
   }
 
+  explicit FP32Vec8(const FP16Vec8& v) {
+    __vector unsigned short raw_u = (__vector unsigned short)v.reg;
+    __vector unsigned int raw_hi =
+        (__vector unsigned int)vec_unpackh((__vector signed short)raw_u);
+    __vector unsigned int raw_lo =
+        (__vector unsigned int)vec_unpackl((__vector signed short)raw_u);
+    reg.val[0] = fp16_to_fp32_bits(raw_hi);
+    reg.val[1] = fp16_to_fp32_bits(raw_lo);
+  }
   float reduce_sum() const {
     AliasReg ar;
     ar.reg = reg;
@@ -410,8 +507,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = vec_xl(48, ptr);
   }
 
+  explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {}
+  explicit FP32Vec16(const FP16Vec16&);
   explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
-
   explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
 
   explicit FP32Vec16(const FP32Vec16& data) {
@@ -435,7 +533,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = data.reg.val[1];
   }
 
-  explicit FP32Vec16(const FP16Vec16& v);
   explicit FP32Vec16(const BF16Vec16& v) {
     reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
     reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -502,28 +599,20 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   FP32Vec16 max(const FP32Vec16& b, int elem_num) const {
     FP32Vec16 result;
 
-    // Create a vector of element indices for each chunk
     __vector unsigned int indices = {0, 1, 2, 3};
     __vector unsigned int elem_num_vec =
         vec_splats(static_cast<unsigned int>(elem_num));
 
-    // Compute masks for each chunk
-    __vector unsigned int chunk_offset0 = {0, 0, 0,
-                                           0};  // Chunk 0: Elements 0-3
-    __vector unsigned int chunk_offset1 = {4, 4, 4,
-                                           4};  // Chunk 1: Elements 4-7
-    __vector unsigned int chunk_offset2 = {8, 8, 8,
-                                           8};  // Chunk 2: Elements 8-11
-    __vector unsigned int chunk_offset3 = {12, 12, 12,
-                                           12};  // Chunk 3: Elements 12-15
+    __vector unsigned int chunk_offset0 = {0, 0, 0, 0};
+    __vector unsigned int chunk_offset1 = {4, 4, 4, 4};
+    __vector unsigned int chunk_offset2 = {8, 8, 8, 8};
+    __vector unsigned int chunk_offset3 = {12, 12, 12, 12};
 
-    // Compute masks for each chunk
     __vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec);
     __vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec);
     __vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec);
     __vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec);
 
-    // Apply masks to compute the result for each chunk
     result.reg.val[0] = vec_sel(this->reg.val[0],
                                 vec_max(this->reg.val[0], b.reg.val[0]), mask0);
     result.reg.val[1] = vec_sel(this->reg.val[1],
@@ -626,6 +715,16 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     vec_xst(reg.val[3], 48, ptr);
   }
 
+  void save(c10::Half* ptr) const {
+    FP16Vec16 fp16_vec(*this);
+    fp16_vec.save(ptr);
+  }
+
+  void save(c10::Half* ptr, const int elem_num) const {
+    FP16Vec16 fp16_vec(*this);
+    fp16_vec.save(ptr, elem_num);
+  }
+
   void save(float* ptr, const int elem_num) const {
     const int elements_in_chunk1 =
         (elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0;
@@ -659,7 +758,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
 };
 
 struct INT8Vec16 : public Vec<INT8Vec16> {
-  constexpr static int VEC_NUM_ELEM = 16;  // 128 bits / 8 bits = 16
+  constexpr static int VEC_NUM_ELEM = 16;
 
   union AliasReg {
     __vector signed char reg;
@@ -707,6 +806,11 @@ struct VecType<c10::BFloat16> {
   using vec_type = BF16Vec8;
 };
 
+template <>
+struct VecType<c10::Half> {
+  using vec_type = FP16Vec8;
+};
+
 template <typename T>
 void storeFP32(float v, T* ptr) {
   *ptr = v;
@@ -723,6 +827,15 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
   *ptr = *(v_ptr + 1);
 }
 
+template <>
+inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
+  __vector float v_vec = {v, 0.0f, 0.0f, 0.0f};
+  __vector unsigned int fp16_bits = fp32_to_fp16_bits(v_vec);
+  unsigned short result =
+      (unsigned short)((__vector unsigned short)fp16_bits)[0];
+  *reinterpret_cast<unsigned short*>(ptr) = result;
+}
+
 #ifndef __VEC_CLASS_FP_NAN
   #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif
@@ -769,38 +882,39 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
 #endif
 }
 
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
+  __vector unsigned int fp16_hi = fp32_to_fp16_bits(v.reg.val[0]);
+  __vector unsigned int fp16_lo = fp32_to_fp16_bits(v.reg.val[1]);
+  reg = (__vector signed short)vec_perm((__vector unsigned char)fp16_hi,
+                                        (__vector unsigned char)fp16_lo, omask);
+}
+
 inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  alignas(16) float temp_fp32[16];
-  alignas(16) c10::Half temp_fp16[16];
-
-  vec_xst(v.reg.val[0], 0, temp_fp32);
-  vec_xst(v.reg.val[1], 16, temp_fp32);
-  vec_xst(v.reg.val[2], 32, temp_fp32);
-  vec_xst(v.reg.val[3], 48, temp_fp32);
-
-  for (int i = 0; i < 16; i++) {
-    temp_fp16[i] = c10::Half(temp_fp32[i]);
-  }
-
-  reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
-  reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
+  __vector unsigned int fp16_0 = fp32_to_fp16_bits(v.reg.val[0]);
+  __vector unsigned int fp16_1 = fp32_to_fp16_bits(v.reg.val[1]);
+  __vector unsigned int fp16_2 = fp32_to_fp16_bits(v.reg.val[2]);
+  __vector unsigned int fp16_3 = fp32_to_fp16_bits(v.reg.val[3]);
+  reg.val[0] = (__vector signed short)vec_perm(
+      (__vector unsigned char)fp16_0, (__vector unsigned char)fp16_1, omask);
+  reg.val[1] = (__vector signed short)vec_perm(
+      (__vector unsigned char)fp16_2, (__vector unsigned char)fp16_3, omask);
 }
 
 inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
-  alignas(16) c10::Half temp_fp16[16];
-  alignas(16) float temp_fp32[16];
-
-  vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
-  vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
-
-  for (int i = 0; i < 16; i++) {
-    temp_fp32[i] = float(temp_fp16[i]);
-  }
-
-  reg.val[0] = vec_xl(0, temp_fp32);
-  reg.val[1] = vec_xl(16, temp_fp32);
-  reg.val[2] = vec_xl(32, temp_fp32);
-  reg.val[3] = vec_xl(48, temp_fp32);
+  __vector unsigned short raw_u0 = (__vector unsigned short)v.reg.val[0];
+  __vector unsigned short raw_u1 = (__vector unsigned short)v.reg.val[1];
+  __vector unsigned int raw_hi0 =
+      (__vector unsigned int)vec_unpackh((__vector signed short)raw_u0);
+  __vector unsigned int raw_lo0 =
+      (__vector unsigned int)vec_unpackl((__vector signed short)raw_u0);
+  __vector unsigned int raw_hi1 =
+      (__vector unsigned int)vec_unpackh((__vector signed short)raw_u1);
+  __vector unsigned int raw_lo1 =
+      (__vector unsigned int)vec_unpackl((__vector signed short)raw_u1);
+  reg.val[0] = fp16_to_fp32_bits(raw_hi0);
+  reg.val[1] = fp16_to_fp32_bits(raw_lo0);
+  reg.val[2] = fp16_to_fp32_bits(raw_hi1);
+  reg.val[3] = fp16_to_fp32_bits(raw_lo1);
 }
 
 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
@@ -864,7 +978,6 @@ inline void prefetch(const void* addr) {
 
 struct INT8Vec64 {
   __vector signed char data[4];
-
   INT8Vec64() = default;
 
   explicit INT8Vec64(const int8_t* ptr) {
@@ -900,5 +1013,4 @@ struct INT8Vec64 {
   void nt_save(int8_t* ptr) const { save(ptr); }
 };
 }  // namespace vec_op
-
 #endif

csrc/cpu/cpu_types_x86.hpp

@@ -3,6 +3,7 @@
 #define CPU_TYPES_X86_HPP
 
 #include <immintrin.h>
+#include <sleef.h>
 #include <torch/all.h>
 
 #ifndef __AVX2__
@@ -592,6 +593,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
 
   FP32Vec16 abs() const { return FP32Vec16(_mm512_abs_ps(reg)); }
 
+  FP32Vec16 tanh() const { return FP32Vec16(Sleef_tanhf16_u10(reg)); }
+
   float reduce_sum() const { return _mm512_reduce_add_ps(reg); }
 
   float reduce_max() const { return _mm512_reduce_max_ps(reg); }
@@ -789,6 +792,12 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                      _mm256_andnot_ps(sign_mask, reg_high));
   }
 
+  FP32Vec16 tanh() const {
+    FP32Vec8 low(reg_low);
+    FP32Vec8 high(reg_high);
+    return FP32Vec16(low.tanh().reg, high.tanh().reg);
+  }
+
   FP32Vec16 min(const FP32Vec16& b) const {
     return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
                      _mm256_min_ps(reg_high, b.reg_high));

csrc/cpu/cpu_wna16.cpp

@@ -155,7 +155,7 @@ void cpu_gemm_wna16_impl(
   constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
   constexpr int32_t n_block_size = 16;
   static_assert(gemm_n_tile_size % n_block_size == 0);
-  const int32_t thread_num = omp_get_max_threads();
+  const int32_t thread_num = cpu_utils::get_max_threads();
 
   // a simple schedule policy, just to hold more B tiles in L2 and make sure
   // each thread has tasks

csrc/cpu/dnnl_kernels.cpp

@@ -202,7 +202,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
   using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
   constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
 
-  const int64_t thread_num = omp_get_max_threads();
+  const int64_t thread_num = cpu_utils::get_max_threads();
   if (num_tokens > thread_num) {
 #pragma omp parallel for
     for (int64_t i = 0; i < num_tokens; ++i) {

csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp

@@ -213,6 +213,8 @@ class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
  public:
   static constexpr int32_t MaxMSize = 32;
   static constexpr int32_t NSize = 32;
+  static constexpr int32_t WeightOCGroupSize = 16;
+  static constexpr bool PackA = false;
 
  public:
   MicroGemm() : curr_m_(-1) {

csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp

@@ -21,6 +21,9 @@ class MicroGemm {
  public:
   static constexpr int32_t MaxMSize = 16;
   static constexpr int32_t NSize = 16;
+  static constexpr int32_t WeightOCGroupSize = 16;
+  // callers must pack A matrix before GEMM
+  static constexpr bool PackA = false;
 
  public:
   void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {

csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp

@@ -0,0 +1,503 @@
+#ifndef CPU_MICRO_GEMM_NEON_HPP
+#define CPU_MICRO_GEMM_NEON_HPP
+
+#include <algorithm>
+#include <cstdint>
+
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+#include <arm_bf16.h>
+#include <arm_neon.h>
+
+namespace cpu_micro_gemm {
+
+namespace {
+
+constexpr int32_t K = 4;
+constexpr int32_t Cols = 2;
+constexpr int32_t TileSize = K * Cols;
+constexpr int32_t Mr = 8;
+constexpr int32_t Nr = 8;
+constexpr int32_t Nr_gemv = 16;
+
+// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1]
+FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) {
+  return vreinterpretq_f32_f64(
+      vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
+}
+
+// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a2, a3, b2, b3]
+FORCE_INLINE float32x4_t zip2_f32x4(const float32x4_t a, const float32x4_t b) {
+  return vreinterpretq_f32_f64(
+      vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
+}
+
+FORCE_INLINE void init_acc_rowpair(float32x4_t& acc01, float32x4_t& acc23,
+                                   float32x4_t& acc45, float32x4_t& acc67,
+                                   const float* __restrict__ c_ptr,
+                                   const int64_t ldc, const int32_t m_rows,
+                                   const bool accum_c) {
+  if (!accum_c || m_rows == 0) {
+    acc01 = vdupq_n_f32(0.0f);
+    acc23 = vdupq_n_f32(0.0f);
+    acc45 = vdupq_n_f32(0.0f);
+    acc67 = vdupq_n_f32(0.0f);
+    return;
+  }
+
+  const float32x4_t row0_0123 = vld1q_f32(c_ptr);
+  const float32x4_t row0_4567 = vld1q_f32(c_ptr + 4);
+  const float32x4_t row1_0123 =
+      (m_rows == 2) ? vld1q_f32(c_ptr + ldc) : vdupq_n_f32(0.0f);
+  const float32x4_t row1_4567 =
+      (m_rows == 2) ? vld1q_f32(c_ptr + ldc + 4) : vdupq_n_f32(0.0f);
+
+  acc01 = zip1_f32x4(row0_0123, row1_0123);
+  acc23 = zip2_f32x4(row0_0123, row1_0123);
+  acc45 = zip1_f32x4(row0_4567, row1_4567);
+  acc67 = zip2_f32x4(row0_4567, row1_4567);
+}
+
+FORCE_INLINE void store_acc_rowpair(const float32x4_t acc01,
+                                    const float32x4_t acc23,
+                                    const float32x4_t acc45,
+                                    const float32x4_t acc67,
+                                    float* __restrict__ c_ptr,
+                                    const int64_t ldc, const int32_t m_rows) {
+  if (m_rows == 0) {
+    return;
+  }
+
+  vst1q_f32(c_ptr, zip1_f32x4(acc01, acc23));
+  vst1q_f32(c_ptr + 4, zip1_f32x4(acc45, acc67));
+
+  if (m_rows == 2) {
+    vst1q_f32(c_ptr + ldc, zip2_f32x4(acc01, acc23));
+    vst1q_f32(c_ptr + ldc + 4, zip2_f32x4(acc45, acc67));
+  }
+}
+
+FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
+    const bfloat16_t* __restrict__ a_packed,
+    const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
+    const int32_t m, const int32_t k_size, const int64_t ldc,
+    const bool accum_c) {
+  float32x4_t acc0101, acc0123, acc0145, acc0167;
+  float32x4_t acc2301, acc2323, acc2345, acc2367;
+  float32x4_t acc4501, acc4523, acc4545, acc4567;
+  float32x4_t acc6701, acc6723, acc6745, acc6767;
+
+  init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
+                   std::min(2, m), accum_c);
+  init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
+                   std::min(2, std::max(0, m - 2)), accum_c);
+  init_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
+                   std::min(2, std::max(0, m - 4)), accum_c);
+  init_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
+                   std::min(2, std::max(0, m - 6)), accum_c);
+
+  const bfloat16_t* __restrict__ a_tile = a_packed;
+  const bfloat16_t* __restrict__ b_tile = b_packed;
+
+#pragma GCC unroll 8
+  for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
+    const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
+    const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
+    const bfloat16x8_t a_tile45 = vld1q_bf16(a_tile + 2 * TileSize);
+    const bfloat16x8_t a_tile67 = vld1q_bf16(a_tile + 3 * TileSize);
+
+    const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile);
+    const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile + TileSize);
+    const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile + 2 * TileSize);
+    const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile + 3 * TileSize);
+
+    acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
+    acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
+    acc4501 = vbfmmlaq_f32(acc4501, a_tile45, b_tile01);
+    acc6701 = vbfmmlaq_f32(acc6701, a_tile67, b_tile01);
+
+    acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
+    acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
+    acc4523 = vbfmmlaq_f32(acc4523, a_tile45, b_tile23);
+    acc6723 = vbfmmlaq_f32(acc6723, a_tile67, b_tile23);
+
+    acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
+    acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
+    acc4545 = vbfmmlaq_f32(acc4545, a_tile45, b_tile45);
+    acc6745 = vbfmmlaq_f32(acc6745, a_tile67, b_tile45);
+
+    acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
+    acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
+    acc4567 = vbfmmlaq_f32(acc4567, a_tile45, b_tile67);
+    acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
+
+    a_tile += 4 * TileSize;
+    b_tile += Nr * K;
+  }
+
+  store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
+                    std::min(2, m));
+  store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
+                    std::min(2, std::max(0, m - 2)));
+  store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
+                    std::min(2, std::max(0, m - 4)));
+  store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
+                    std::min(2, std::max(0, m - 6)));
+}
+
+FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
+    const bfloat16_t* __restrict__ a_packed,
+    const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
+    const int32_t m, const int32_t k_size, const int64_t b_n_group_stride,
+    const int64_t ldc, const bool accum_c) {
+  const int32_t m_rows_01 = std::min(2, m);
+  const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
+
+  float32x4_t acc0101, acc0123, acc0145, acc0167;
+  float32x4_t acc2301, acc2323, acc2345, acc2367;
+  float32x4_t acc0189, acc011011, acc011213, acc011415;
+  float32x4_t acc2389, acc231011, acc231213, acc231415;
+
+  init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01,
+                   accum_c);
+  init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
+                   m_rows_23, accum_c);
+  init_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
+                   m_rows_01, accum_c);
+  init_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
+                   c_ptr + 2 * ldc + 8, ldc, m_rows_23, accum_c);
+
+  const bfloat16_t* __restrict__ a_tile = a_packed;
+  const bfloat16_t* __restrict__ b_tile0 = b_packed;
+  const bfloat16_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
+
+#pragma GCC unroll 8
+  for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
+    const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
+    const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
+    const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile0);
+    const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile0 + TileSize);
+    const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile0 + 2 * TileSize);
+    const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile0 + 3 * TileSize);
+    const bfloat16x8_t b_tile89 = vld1q_bf16(b_tile1);
+    const bfloat16x8_t b_tile1011 = vld1q_bf16(b_tile1 + TileSize);
+    const bfloat16x8_t b_tile1213 = vld1q_bf16(b_tile1 + 2 * TileSize);
+    const bfloat16x8_t b_tile1415 = vld1q_bf16(b_tile1 + 3 * TileSize);
+
+    acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
+    acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
+    acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
+    acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
+
+    acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
+    acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
+    acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
+    acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
+
+    acc0189 = vbfmmlaq_f32(acc0189, a_tile01, b_tile89);
+    acc2389 = vbfmmlaq_f32(acc2389, a_tile23, b_tile89);
+    acc011011 = vbfmmlaq_f32(acc011011, a_tile01, b_tile1011);
+    acc231011 = vbfmmlaq_f32(acc231011, a_tile23, b_tile1011);
+
+    acc011213 = vbfmmlaq_f32(acc011213, a_tile01, b_tile1213);
+    acc231213 = vbfmmlaq_f32(acc231213, a_tile23, b_tile1213);
+    acc011415 = vbfmmlaq_f32(acc011415, a_tile01, b_tile1415);
+    acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);
+
+    a_tile += 2 * TileSize;
+    b_tile0 += Nr * K;
+    b_tile1 += Nr * K;
+  }
+
+  store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
+  store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
+                    m_rows_23);
+  store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
+                    m_rows_01);
+  store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
+                    c_ptr + 2 * ldc + 8, ldc, m_rows_23);
+}
+
+}  // namespace
+
+template <typename scalar_t>
+class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
+ public:
+  static constexpr int32_t MaxMSize = 8;
+  static constexpr int32_t NSize = 32;
+  static constexpr int32_t WeightOCGroupSize = Nr;
+  static constexpr bool PackA = false;
+
+ public:
+  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
+  }
+
+  static void pack_weight(const scalar_t* __restrict__ /*weight*/,
+                          scalar_t* __restrict__ /*packed_weight*/,
+                          const int32_t /*output_size*/,
+                          const int32_t /*input_size*/) {
+    TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
+  }
+};
+
+template <>
+class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
+ public:
+  using scalar_t = c10::BFloat16;
+
+  static constexpr int32_t MaxMSize = 8;
+  static constexpr int32_t NSize = 32;
+  static constexpr int32_t WeightOCGroupSize = Nr;
+  static constexpr bool PackA = true;
+
+ public:
+  // physical layout [
+  //  M / 8; Mr is 8
+  //  K / 4; K for bfmmla is 4
+  //  4,   ; 4 row-pairs for each 8 rows
+  //  2,   ; row-pair is 2 rows
+  //  4    ; 4 elements per row
+  // ]
+
+  static void pack_input_from_rows(const scalar_t* const* __restrict__ rows,
+                                   scalar_t* __restrict__ a_packed,
+                                   const int32_t m, const int32_t k) {
+    TORCH_CHECK(m > 0 && m <= MaxMSize);
+    TORCH_CHECK_EQ(k % K, 0);
+
+    auto* __restrict__ out = reinterpret_cast<bfloat16_t*>(a_packed);
+    const bfloat16x8_t zero_q = vdupq_n_bf16(bfloat16_t{});
+    const bfloat16x4_t zero = vget_low_bf16(zero_q);
+
+    for (int32_t row_base = 0; row_base < m; row_base += Mr) {
+      const int32_t actual_m = std::min(Mr, m - row_base);
+      const bfloat16_t* __restrict__ row[Mr];
+      for (int32_t i = 0; i < actual_m; ++i) {
+        row[i] = reinterpret_cast<const bfloat16_t*>(rows[row_base + i]);
+      }
+
+      if (actual_m == 8) {
+        int32_t k_idx = 0;
+        for (; k_idx + 8 <= k; k_idx += 8) {
+          bfloat16_t* __restrict__ block0 = out;
+          bfloat16_t* __restrict__ block1 = out + 4 * TileSize;
+
+          bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
+          bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
+          vst1q_bf16(block0,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          a0 = vld1q_bf16(row[2] + k_idx);
+          a1 = vld1q_bf16(row[3] + k_idx);
+          vst1q_bf16(block0 + TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          a0 = vld1q_bf16(row[4] + k_idx);
+          a1 = vld1q_bf16(row[5] + k_idx);
+          vst1q_bf16(block0 + 2 * TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + 2 * TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          a0 = vld1q_bf16(row[6] + k_idx);
+          a1 = vld1q_bf16(row[7] + k_idx);
+          vst1q_bf16(block0 + 3 * TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + 3 * TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          out += 8 * TileSize;
+        }
+
+        for (; k_idx < k; k_idx += K) {
+          bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
+          bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
+          vst1q_bf16(out, vcombine_bf16(a0, a1));
+
+          a0 = vld1_bf16(row[2] + k_idx);
+          a1 = vld1_bf16(row[3] + k_idx);
+          vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
+
+          a0 = vld1_bf16(row[4] + k_idx);
+          a1 = vld1_bf16(row[5] + k_idx);
+          vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
+
+          a0 = vld1_bf16(row[6] + k_idx);
+          a1 = vld1_bf16(row[7] + k_idx);
+          vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
+
+          out += 4 * TileSize;
+        }
+        continue;
+      }
+
+      if (actual_m == 4) {
+        int32_t k_idx = 0;
+        for (; k_idx + 8 <= k; k_idx += 8) {
+          bfloat16_t* __restrict__ block0 = out;
+          bfloat16_t* __restrict__ block1 = out + 2 * TileSize;
+
+          bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
+          bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
+          vst1q_bf16(block0,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          a0 = vld1q_bf16(row[2] + k_idx);
+          a1 = vld1q_bf16(row[3] + k_idx);
+          vst1q_bf16(block0 + TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          out += 4 * TileSize;
+        }
+
+        for (; k_idx < k; k_idx += K) {
+          bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
+          bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
+          vst1q_bf16(out, vcombine_bf16(a0, a1));
+
+          a0 = vld1_bf16(row[2] + k_idx);
+          a1 = vld1_bf16(row[3] + k_idx);
+          vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
+
+          out += 2 * TileSize;
+        }
+        continue;
+      }
+
+      const int32_t row_pair_count = (actual_m <= 4) ? 2 : Mr / 2;
+
+      int32_t k_idx = 0;
+      for (; k_idx + 8 <= k; k_idx += 8) {
+        bfloat16_t* __restrict__ block0 = out;
+        bfloat16_t* __restrict__ block1 = out + row_pair_count * TileSize;
+
+        bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
+        bfloat16x8_t a1 = (actual_m > 1) ? vld1q_bf16(row[1] + k_idx) : zero_q;
+        vst1q_bf16(block0, vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+        vst1q_bf16(block1,
+                   vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+        a0 = (actual_m > 2) ? vld1q_bf16(row[2] + k_idx) : zero_q;
+        a1 = (actual_m > 3) ? vld1q_bf16(row[3] + k_idx) : zero_q;
+        vst1q_bf16(block0 + TileSize,
+                   vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+        vst1q_bf16(block1 + TileSize,
+                   vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+        if (actual_m > 4) {
+          a0 = vld1q_bf16(row[4] + k_idx);
+          a1 = (actual_m > 5) ? vld1q_bf16(row[5] + k_idx) : zero_q;
+          vst1q_bf16(block0 + 2 * TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + 2 * TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+
+          a0 = (actual_m > 6) ? vld1q_bf16(row[6] + k_idx) : zero_q;
+          a1 = (actual_m > 7) ? vld1q_bf16(row[7] + k_idx) : zero_q;
+          vst1q_bf16(block0 + 3 * TileSize,
+                     vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
+          vst1q_bf16(block1 + 3 * TileSize,
+                     vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
+        }
+
+        out += 2 * row_pair_count * TileSize;
+      }
+
+      for (; k_idx < k; k_idx += K) {
+        bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
+        bfloat16x4_t a1 = (actual_m > 1) ? vld1_bf16(row[1] + k_idx) : zero;
+        vst1q_bf16(out, vcombine_bf16(a0, a1));
+
+        a0 = (actual_m > 2) ? vld1_bf16(row[2] + k_idx) : zero;
+        a1 = (actual_m > 3) ? vld1_bf16(row[3] + k_idx) : zero;
+        vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
+
+        if (actual_m > 4) {
+          a0 = vld1_bf16(row[4] + k_idx);
+          a1 = (actual_m > 5) ? vld1_bf16(row[5] + k_idx) : zero;
+          vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
+
+          a0 = (actual_m > 6) ? vld1_bf16(row[6] + k_idx) : zero;
+          a1 = (actual_m > 7) ? vld1_bf16(row[7] + k_idx) : zero;
+          vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
+        }
+        out += row_pair_count * TileSize;
+      }
+    }
+  }
+
+  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    (void)lda;  // A is packed, so lda is not needed
+    TORCH_CHECK_EQ(k % K, 0);
+
+    for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
+      const bfloat16_t* __restrict__ b_panel =
+          reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
+
+      for (int32_t row_base = 0; row_base < m; row_base += Mr) {
+        const int32_t panel_m = std::min(Mr, m - row_base);
+        const bfloat16_t* __restrict__ a_panel =
+            reinterpret_cast<const bfloat16_t*>(a_ptr) + row_base * k;
+        float* __restrict__ c_panel = c_ptr + row_base * ldc + n_idx;
+
+        if (panel_m <= 4) {
+          gemm_micro_bfmmla_4x16_packed_a(a_panel, b_panel, c_panel, panel_m, k,
+                                          b_n_group_stride, ldc, accum_c);
+        } else {
+          gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel, c_panel, panel_m, k,
+                                         ldc, accum_c);
+          gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel + b_n_group_stride,
+                                         c_panel + Nr, panel_m, k, ldc,
+                                         accum_c);
+        }
+      }
+    }
+  }
+
+  // physical layout [
+  //  N / 8; Nr is 8
+  //  K / 4; K for bfmmla is 4
+  //  4,   ; 4 col-pairs for each 8 cols
+  //  2,   ; col-pair is 2 cols
+  //  4    ; 4 elements per col
+  // ]
+  static void pack_weight(const c10::BFloat16* __restrict__ weight,
+                          c10::BFloat16* __restrict__ packed_weight,
+                          const int32_t output_size, const int32_t input_size) {
+    TORCH_CHECK_EQ(output_size % NSize, 0);
+    TORCH_CHECK_EQ(input_size % K, 0);
+
+    for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
+      c10::BFloat16* __restrict__ dst = packed_weight + o_idx * input_size;
+      for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
+        for (int32_t pair_idx = 0; pair_idx < Nr; pair_idx += Cols) {
+          const c10::BFloat16* __restrict__ row0 =
+              weight + (o_idx + pair_idx) * input_size;
+          const c10::BFloat16* __restrict__ row1 = row0 + input_size;
+          dst[0] = row0[k_idx + 0];
+          dst[1] = row0[k_idx + 1];
+          dst[2] = row0[k_idx + 2];
+          dst[3] = row0[k_idx + 3];
+          dst[4] = row1[k_idx + 0];
+          dst[5] = row1[k_idx + 1];
+          dst[6] = row1[k_idx + 2];
+          dst[7] = row1[k_idx + 3];
+          dst += TileSize;
+        }
+      }
+    }
+  }
+};
+
+}  // namespace cpu_micro_gemm
+
+#endif

csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp

@@ -104,6 +104,8 @@ class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
  public:
   static constexpr int32_t MaxMSize = 8;
   static constexpr int32_t NSize = 32;
+  static constexpr int32_t WeightOCGroupSize = 16;
+  static constexpr bool PackA = false;
 
  public:
   void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {

csrc/cpu/mla_decode.cpp

@@ -18,17 +18,9 @@ struct KernelVecType<float> {
 
 template <>
 struct KernelVecType<c10::Half> {
-#if defined(__powerpc64__)
-  // Power specific vector types
-  using qk_load_vec_type = vec_op::FP32Vec16;
-  using qk_vec_type = vec_op::FP32Vec16;
-  using v_load_vec_type = vec_op::FP32Vec16;
-#else
-  // Fallback for other architectures, including x86
   using qk_load_vec_type = vec_op::FP16Vec16;
   using qk_vec_type = vec_op::FP32Vec16;
   using v_load_vec_type = vec_op::FP16Vec16;
-#endif
 };
 
 #ifdef __AVX512BF16__
@@ -259,7 +251,7 @@ void mla_decode_kvcache_cpu_impl(
   constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
 
   // shared across threads
-  const int max_threads = omp_get_max_threads();
+  const int max_threads = cpu_utils::get_max_threads();
   const int acc_out_nbytes =
       max_threads * num_heads * V_HEAD_DIM * sizeof(float);
   float* acc_out = static_cast<float*>(std::aligned_alloc(64, acc_out_nbytes));

csrc/cpu/pos_encoding.cpp

@@ -1,4 +1,3 @@
-
 #include "cpu_types.hpp"
 
 namespace {
@@ -97,6 +96,91 @@ void rotary_embedding_impl(
   }
 }
 
+template <>
+void rotary_embedding_impl<c10::Half>(
+    const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
+    c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
+    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
+    const int num_heads, const int num_kv_heads, const int head_size,
+    const int num_tokens) {
+  using scalar_vec_t = vec_op::FP16Vec8;
+  constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
+
+  const int embed_dim = rot_dim / 2;
+  bool flag = (embed_dim % VEC_ELEM_NUM == 0);
+  const int loop_upper = flag ? embed_dim : embed_dim - VEC_ELEM_NUM;
+
+  auto compute_loop = [&](const int64_t token_head, const c10::Half* cache_ptr,
+                          c10::Half* qk) {
+    int j = 0;
+    for (; j < loop_upper; j += VEC_ELEM_NUM) {
+      const int rot_offset = j;
+      const int x_index = rot_offset;
+      const int y_index = embed_dim + rot_offset;
+
+      const int64_t out_x = token_head + x_index;
+      const int64_t out_y = token_head + y_index;
+
+      const vec_op::FP16Vec8 cos_fp16(cache_ptr + x_index);
+      const vec_op::FP16Vec8 sin_fp16(cache_ptr + y_index);
+      const vec_op::FP16Vec8 q_x_fp16(qk + out_x);
+      const vec_op::FP16Vec8 q_y_fp16(qk + out_y);
+
+      const vec_op::FP32Vec8 fp32_cos(cos_fp16);
+      const vec_op::FP32Vec8 fp32_sin(sin_fp16);
+      const vec_op::FP32Vec8 fp32_q_x(q_x_fp16);
+      const vec_op::FP32Vec8 fp32_q_y(q_y_fp16);
+
+      auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin;
+      auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin;
+
+      vec_op::FP16Vec8(out1).save(qk + out_x);
+      vec_op::FP16Vec8(out2).save(qk + out_y);
+    }
+    if (!flag) {
+      for (; j < embed_dim; ++j) {
+        const int x_index = j;
+        const int y_index = embed_dim + j;
+
+        const int64_t out_x = token_head + x_index;
+        const int64_t out_y = token_head + y_index;
+
+        const float fp32_cos = static_cast<float>(cache_ptr[x_index]);
+        const float fp32_sin = static_cast<float>(cache_ptr[y_index]);
+        const float fp32_q_x = static_cast<float>(qk[out_x]);
+        const float fp32_q_y = static_cast<float>(qk[out_y]);
+
+        qk[out_x] =
+            static_cast<c10::Half>(fp32_q_x * fp32_cos - fp32_q_y * fp32_sin);
+        qk[out_y] =
+            static_cast<c10::Half>(fp32_q_y * fp32_cos + fp32_q_x * fp32_sin);
+      }
+    }
+  };
+
+#pragma omp parallel for
+  for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
+    int64_t pos = positions[token_idx];
+    const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
+
+    for (int i = 0; i < num_heads; ++i) {
+      const int head_idx = i;
+      const int64_t token_head =
+          token_idx * query_stride + head_idx * head_size;
+      compute_loop(token_head, cache_ptr, query);
+    }
+
+    if (key != nullptr) {
+      for (int i = 0; i < num_kv_heads; ++i) {
+        const int head_idx = i;
+        const int64_t token_head =
+            token_idx * key_stride + head_idx * head_size;
+        compute_loop(token_head, cache_ptr, key);
+      }
+    }
+  }
+}
+
 template <typename scalar_t>
 void rotary_embedding_gptj_impl(
     const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
@@ -174,6 +258,75 @@ void rotary_embedding_gptj_impl(
     }
   }
 }
+
+template <>
+void rotary_embedding_gptj_impl<c10::Half>(
+    const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
+    c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
+    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
+    const int num_heads, const int num_kv_heads, const int head_size,
+    const int num_tokens) {
+  const int embed_dim = rot_dim / 2;
+
+#pragma omp parallel for collapse(2)
+  for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
+    for (int i = 0; i < num_heads; ++i) {
+      int64_t pos = positions[token_idx];
+      const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
+      const c10::Half* cos_cache_ptr = cache_ptr;
+      const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
+      const int head_idx = i;
+      const int64_t token_head =
+          token_idx * query_stride + head_idx * head_size;
+      c10::Half* head_query = token_head + query;
+      for (int j = 0; j < embed_dim; j += 1) {
+        const int rot_offset = j;
+        const int x_index = 2 * rot_offset;
+        const int y_index = 2 * rot_offset + 1;
+
+        const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
+        const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
+
+        const float x = static_cast<float>(head_query[x_index]);
+        const float y = static_cast<float>(head_query[y_index]);
+
+        head_query[x_index] = static_cast<c10::Half>(x * cos - y * sin);
+        head_query[y_index] = static_cast<c10::Half>(y * cos + x * sin);
+      }
+    }
+  }
+
+  if (key == nullptr) {
+    return;
+  }
+
+#pragma omp parallel for collapse(2)
+  for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
+    for (int i = 0; i < num_kv_heads; ++i) {
+      int64_t pos = positions[token_idx];
+      const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
+      const c10::Half* cos_cache_ptr = cache_ptr;
+      const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
+      const int head_idx = i;
+      const int64_t token_head = token_idx * key_stride + head_idx * head_size;
+      c10::Half* head_key = key + token_head;
+      for (int j = 0; j < embed_dim; j += 1) {
+        const int rot_offset = j;
+        const int x_index = 2 * rot_offset;
+        const int y_index = 2 * rot_offset + 1;
+
+        const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
+        const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
+
+        const float x = static_cast<float>(head_key[x_index]);
+        const float y = static_cast<float>(head_key[y_index]);
+
+        head_key[x_index] = static_cast<c10::Half>(x * cos - y * sin);
+        head_key[y_index] = static_cast<c10::Half>(y * cos + x * sin);
+      }
+    }
+  }
+}
 };  // namespace
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,

csrc/cpu/sgl-kernels/conv.cpp

@@ -289,19 +289,18 @@ void causal_conv1d_fwd_kernel_impl(
   }
 }
 
-#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE)                   \
-  tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
-      input + batch_offset * dim + mb_start * dim + nb_start,       \
-      weight + nb_start * width,                                    \
-      out + batch_offset * dim + mb_start * dim + nb_start,         \
-      has_bias ? bias + nb_start : nullptr,                         \
-      nullptr,                                                      \
-      false,                                                        \
-      mb_size,                                                      \
-      dim,                                                          \
+#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE)                                                   \
+  tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply(                                 \
+      input + batch_offset * dim + mb_start * dim + nb_start,                                       \
+      weight + nb_start * width,                                                                    \
+      out + batch_offset * dim + mb_start * dim + nb_start,                                         \
+      has_bias ? bias + nb_start : nullptr,                                                         \
+      has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
+      has_initial_states_value,                                                                     \
+      mb_size,                                                                                      \
+      dim,                                                                                          \
       mb_start == 0);
 
-// TODO: add `has_initial_state` support for varlen kernel
 template <typename scalar_t>
 void causal_conv1d_fwd_varlen_kernel_impl(
     scalar_t* __restrict__ out,
@@ -343,6 +342,9 @@ void causal_conv1d_fwd_varlen_kernel_impl(
         int64_t nb_start = nb * BLOCK_N;
         int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
 
+        const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
+        int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
+
         switch (width << 4 | nb_size >> 4) {
           case 0x42:
             LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
@@ -373,7 +375,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
             width,
             dim,
             seqlen,
-            /* has_initial_state */ false);
+            has_initial_state[bs]);
       }
     });
   }

csrc/cpu/sgl-kernels/gemm_int4.cpp

@@ -285,6 +285,125 @@ inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t
   return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
 }
 
+#if defined(CPU_CAPABILITY_RVV)
+template <int64_t N, int64_t ldb, int group>
+inline fixed_i8x8_t load_uint4_as_int8_rvv(const uint8_t* __restrict__ B, int64_t k) {
+  constexpr int64_t n_group_size = 8;
+  constexpr int64_t vnni_size = 4;
+  static_assert(N == 32);
+  static_assert(ldb == N / 2);
+  static_assert(group >= 0 && group < N / n_group_size);
+
+  // Unpack: gather 8 packed int4 values from the VNNI4 layout.
+  const int64_t ki = k % vnni_size;
+  const int64_t k_base = k - ki;
+  constexpr int64_t packed_group = group / 2;
+  const uint8_t* packed_ptr = B + k_base * ldb + packed_group * n_group_size * vnni_size + ki;
+
+  fixed_u8x8_t packed = RVVI(__riscv_vlse8_v_u8, LMUL_64)(packed_ptr, vnni_size, n_group_size);
+  if constexpr (group % 2 == 1) {
+    packed = RVVI(__riscv_vsrl_vx_u8, LMUL_64)(packed, 4, n_group_size);
+  }
+  fixed_u8x8_t nibbles = RVVI(__riscv_vand_vx_u8, LMUL_64)(packed, 0x0f, n_group_size);
+  return RVVI4(__riscv_vreinterpret_v_u8, LMUL_64, _i8, LMUL_64)(nibbles);
+}
+
+inline fixed_i32x8_t gemm_accum_uint8_int8_rvv(fixed_i32x8_t acc, uint8_t a, fixed_i8x8_t b) {
+  constexpr int64_t vl = 8;
+  fixed_i16x8_t b_i16 = RVVI(__riscv_vsext_vf2_i16, LMUL_128)(b, vl);
+  return RVVI(__riscv_vwmacc_vx_i32, LMUL_256)(acc, static_cast<int16_t>(a), b_i16, vl);
+}
+
+template <int64_t N, int64_t ldb, int group>
+inline fixed_i32x8_t gemm_accum_uint4_rvv(
+    fixed_i32x8_t acc,
+    const uint8_t* __restrict__ B,
+    const int8_t* __restrict__ qzeros_b,
+    uint8_t a,
+    int64_t k) {
+  constexpr int64_t n_group_size = 8;
+  fixed_i8x8_t b = load_uint4_as_int8_rvv<N, ldb, group>(B, k);
+  fixed_i8x8_t qzeros =
+      RVVI(__riscv_vle8_v_i8, LMUL_64)(qzeros_b + group * n_group_size, n_group_size);
+  b = RVVI(__riscv_vsub_vv_i8, LMUL_64)(b, qzeros, n_group_size);
+  return gemm_accum_uint8_int8_rvv(acc, a, b);
+}
+
+template <int group>
+inline void _dequant_and_store_rvv(
+    float* __restrict__ C,
+    fixed_i32x8_t acc,
+    const float* __restrict__ scales_a,
+    const int32_t* __restrict__ qzeros_a,
+    const float* __restrict__ scales_b,
+    const int32_t* __restrict__ compensation,
+    int64_t m,
+    int64_t ldc) {
+  constexpr int64_t n_group_size = 8;
+  constexpr int64_t n = group * n_group_size;
+  constexpr int64_t vl = n_group_size;
+
+  // Dequant compensation: remove activation zero-point contribution.
+  fixed_i32x8_t comp = RVVI(__riscv_vle32_v_i32, LMUL_256)(compensation + n, vl);
+  fixed_i32x8_t zp_comp = RVVI(__riscv_vmul_vx_i32, LMUL_256)(comp, qzeros_a[m], vl);
+  acc = RVVI(__riscv_vsub_vv_i32, LMUL_256)(acc, zp_comp, vl);
+
+  // Scale: convert int32 accumulators to fp32 and apply activation/weight scales.
+  fixed_fp32x8_t acc_f = RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(acc, vl);
+  acc_f = RVVI(__riscv_vfmul_vf_f32, LMUL_256)(acc_f, scales_a[m], vl);
+  fixed_fp32x8_t scale_b = RVVI(__riscv_vle32_v_f32, LMUL_256)(scales_b + n, vl);
+  acc_f = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(acc_f, scale_b, vl);
+
+  // Store: accumulate into the float scratch buffer that already holds bias/zero.
+  float* c_ptr = C + m * ldc + n;
+  fixed_fp32x8_t c_old = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr, vl);
+  fixed_fp32x8_t c_new = RVVI(__riscv_vfadd_vv_f32, LMUL_256)(c_old, acc_f, vl);
+  RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr, c_new, vl);
+}
+
+template <int64_t N, int64_t ldb>
+void _dequant_gemm_accum_rvv(
+    float* __restrict__ C,
+    const uint8_t* __restrict__ A,
+    const float* __restrict__ scales_a,
+    const int32_t* __restrict__ qzeros_a,
+    const uint8_t* __restrict__ B,
+    const float* __restrict__ scales_b,
+    const int8_t* __restrict__ qzeros_b,
+    const int32_t* __restrict__ compensation,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc) {
+  static_assert(N == 32);
+  static_assert(ldb == N / 2);
+  constexpr int64_t vl = 8;
+
+  // Accumulate one C row over the 32-column block.
+  for (int64_t m = 0; m < M; ++m) {
+    fixed_i32x8_t acc0 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
+    fixed_i32x8_t acc1 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
+    fixed_i32x8_t acc2 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
+    fixed_i32x8_t acc3 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
+    // A[m][k] @ B[k][0:32] -> acc[m][0:32]
+    for (int64_t k = 0; k < K; ++k) {
+      // GEMM K step: one scalar activation updates four 8-column RVV tiles.
+      const uint8_t a = A[m * lda + k];
+      acc0 = gemm_accum_uint4_rvv<N, ldb, 0>(acc0, B, qzeros_b, a, k);
+      acc1 = gemm_accum_uint4_rvv<N, ldb, 1>(acc1, B, qzeros_b, a, k);
+      acc2 = gemm_accum_uint4_rvv<N, ldb, 2>(acc2, B, qzeros_b, a, k);
+      acc3 = gemm_accum_uint4_rvv<N, ldb, 3>(acc3, B, qzeros_b, a, k);
+    }
+
+    // Dequant/scale/store each 8-column group back into C.
+    _dequant_and_store_rvv<0>(C, acc0, scales_a, qzeros_a, scales_b, compensation, m, ldc);
+    _dequant_and_store_rvv<1>(C, acc1, scales_a, qzeros_a, scales_b, compensation, m, ldc);
+    _dequant_and_store_rvv<2>(C, acc2, scales_a, qzeros_a, scales_b, compensation, m, ldc);
+    _dequant_and_store_rvv<3>(C, acc3, scales_a, qzeros_a, scales_b, compensation, m, ldc);
+  }
+}
+#endif
+
 template <int64_t N, int64_t ldb, bool sym_quant_act>
 void _dequant_gemm_accum(
     float* C,
@@ -336,6 +455,11 @@ void _dequant_gemm_accum(
     _dequant_and_store<true, N, sym_quant_act>(
         C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
   } else
+#elif defined(CPU_CAPABILITY_RVV)
+  if constexpr (!sym_quant_act && N == BLOCK_N && ldb == BLOCK_N / 2) {
+    _dequant_gemm_accum_rvv<N, ldb>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, M, K, lda, ldc);
+    return;
+  } else
 #endif
   {
     for (int64_t m = 0; m < M; ++m) {

csrc/cpu/sgl-kernels/vec.h

@@ -9,11 +9,19 @@
 #define CPU_CAPABILITY_AVX512
 #endif
 
+#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
+#define CPU_CAPABILITY_RVV
+#endif
+
 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
 #if defined(CPU_CAPABILITY_AVX512)
 #include <immintrin.h>
 #endif
+
+#if defined(CPU_CAPABILITY_RVV)
+#include "../cpu_types_riscv_defs.hpp"
+#endif
 namespace {
 
 using namespace at::vec;

csrc/cpu/spec_decode_utils.cpp

@@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl(
   }
 }
 
+void copy_and_expand_dflash_inputs_kernel_impl(
+    const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
+    torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
+    torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
+    torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
+    const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
+    const std::optional<torch::Tensor>& num_rejected_tokens,
+    const int64_t parallel_drafting_token_id, const int64_t block_size,
+    const int64_t num_query_per_req, const int64_t num_speculative_tokens,
+    const int64_t total_input_tokens, const bool has_num_rejected) {
+  const int64_t num_reqs = query_start_loc.size(0) - 1;
+
+  const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
+  const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
+  const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
+  const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
+  const int64_t* rejected_ptr =
+      has_num_rejected && num_rejected_tokens.has_value()
+          ? num_rejected_tokens.value().data_ptr<int64_t>()
+          : nullptr;
+
+  int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
+  int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>();
+  int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>();
+  int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>();
+  int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>();
+  int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>();
+
+  const int64_t block_table_stride = block_table.stride(0);
+
+#pragma omp parallel for
+  for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
+    int32_t ctx_start = query_start_ptr[req_idx];
+    int32_t ctx_end = query_start_ptr[req_idx + 1];
+    int64_t num_ctx = ctx_end - ctx_start;
+    int64_t valid_ctx_end = ctx_end;
+    if (rejected_ptr != nullptr) {
+      valid_ctx_end -= rejected_ptr[req_idx];
+    }
+    // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that
+    // valid_ctx_end - 1 never reads before the request's context range.
+    valid_ctx_end =
+        std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1));
+
+    int64_t last_pos = target_pos_ptr[valid_ctx_end - 1];
+
+    for (int64_t j = 0; j < num_ctx; ++j) {
+      int64_t ctx_idx = ctx_start + j;
+      int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1);
+      int64_t position = target_pos_ptr[ctx_pos_idx];
+      int64_t block_num = position / block_size;
+      block_num = std::min(block_num, block_table_stride - 1);
+      int32_t block_id =
+          block_table_ptr[req_idx * block_table_stride + block_num];
+      int64_t slot = block_id * block_size + (position % block_size);
+
+      out_ctx_pos_ptr[ctx_idx] = position;
+      out_ctx_slot_ptr[ctx_idx] = slot;
+    }
+
+    for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) {
+      int64_t query_out = req_idx * num_query_per_req + query_off;
+      int64_t position = last_pos + 1 + query_off;
+      int64_t block_num = position / block_size;
+      block_num = std::min(block_num, block_table_stride - 1);
+      int32_t block_id =
+          block_table_ptr[req_idx * block_table_stride + block_num];
+      int64_t slot = block_id * block_size + (position % block_size);
+
+      out_query_pos_ptr[query_out] = position;
+      out_query_slot_ptr[query_out] = slot;
+      out_ids_ptr[query_out] =
+          query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id;
+
+      if (query_off > 0) {
+        int64_t sample_out_idx =
+            req_idx * num_speculative_tokens + (query_off - 1);
+        out_token_idx_ptr[sample_out_idx] = query_out;
+      }
+    }
+  }
+}
+
 void rejection_greedy_sample_kernel_impl(
     torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,

csrc/cpu/torch_bindings.cpp

@@ -237,6 +237,16 @@ void copy_and_expand_eagle_inputs_kernel_impl(
     const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
     const int64_t total_input_tokens,
     const int64_t num_padding_slots_per_request, const bool shift_input_ids);
+void copy_and_expand_dflash_inputs_kernel_impl(
+    const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
+    torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
+    torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
+    torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
+    const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
+    const std::optional<torch::Tensor>& num_rejected_tokens,
+    const int64_t parallel_drafting_token_id, const int64_t block_size,
+    const int64_t num_query_per_req, const int64_t num_speculative_tokens,
+    const int64_t total_input_tokens, const bool has_num_rejected);
 void rejection_greedy_sample_kernel_impl(
     torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
@@ -288,6 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
 
+  // GELU tanh implementation.
+  ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
+  ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
+
   // GELU implementation used in GPT-2.
   ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_new", torch::kCPU, &gelu_new);
@@ -538,7 +552,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif
 
   // fused moe
-#if defined(__AVX512F__)
+#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
   ops.def(
       "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
       "-> ()");
@@ -549,7 +563,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "bool skip_weighted, "
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
-#endif
+#endif  // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
   ops.def(
       "mla_decode_kvcache("
       "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -599,6 +613,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
       "bool shift_input_ids) -> ()",
       &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
+  ops.def(
+      "copy_and_expand_dflash_inputs_kernel_impl("
+      "Tensor next_token_ids, Tensor target_positions, "
+      "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, "
+      "Tensor(a4!) out_query_positions, "
+      "Tensor(a5!) out_context_slot_mapping, "
+      "Tensor(a6!) out_query_slot_mapping, "
+      "Tensor(a7!) out_token_indices, Tensor block_table, "
+      "Tensor query_start_loc, Tensor? num_rejected_tokens, "
+      "SymInt parallel_drafting_token_id, SymInt block_size, "
+      "SymInt num_query_per_req, SymInt num_speculative_tokens, "
+      "SymInt total_input_tokens, bool has_num_rejected) -> ()",
+      &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl);
   ops.def(
       "rejection_greedy_sample_kernel_impl("
       "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "

csrc/cpu/utils.hpp

@@ -2,13 +2,14 @@
 #define UTILS_HPP
 
 #include <atomic>
+#include <string>
 #include <unistd.h>
 #include <ATen/cpu/Utils.h>
 
 #include "cpu/cpu_types.hpp"
 
 namespace cpu_utils {
-enum class ISA { AMX, VEC, RVV };
+enum class ISA { AMX, VEC, RVV, NEON };
 
 inline ISA get_isa(const std::string& isa) {
   if (isa == "amx") {
@@ -17,6 +18,8 @@ inline ISA get_isa(const std::string& isa) {
     return ISA::VEC;
   } else if (isa == "rvv") {
     return ISA::RVV;
+  } else if (isa == "neon") {
+    return ISA::NEON;
   } else {
     TORCH_CHECK(false, "Invalid isa type: " + isa);
   }

csrc/cuda_view.cu

@@ -1,60 +0,0 @@
-// TODO: Remove this once ROCm upgrade to torch 2.11.
-#include <torch/all.h>
-#include <torch/cuda.h>
-#include <cuda_runtime.h>
-
-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
-  TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
-
-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::empty(cpu_tensor.sizes(),
-                        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  if (cpu_tensor.is_pinned()) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    TORCH_CHECK(err == cudaSuccess,
-                "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-
-    return torch::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
-        [base = cpu_tensor](void*) {},  // keep cpu tensor alive
-        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
-  size_t nbytes = contiguous_cpu.nbytes();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
-
-  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
-
-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
-
-  return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
-                          contiguous_cpu.strides(), deleter,
-                          contiguous_cpu.options().device(torch::kCUDA));
-}

csrc/cumem_allocator.cpp

@@ -48,8 +48,8 @@ static inline unsigned long long my_min(unsigned long long a,
 }
 
 static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
-                                     size_t alignment) {
-  CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);
+                                     size_t alignment, CUdeviceptr addr = 0) {
+  CUresult status = cuMemAddressReserve(d_mem, size, alignment, addr, 0);
   if (status == CUresult(0) || alignment == 0) {
     return status;
   }
@@ -58,7 +58,7 @@ static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
   // alignment even when physical VRAM is free. Let HIP choose the default
   // alignment, then verify that the returned address still satisfies the
   // requested alignment before accepting it.
-  status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
+  status = cuMemAddressReserve(d_mem, size, 0, addr, 0);
   if (status != CUresult(0)) {
     return status;
   }
@@ -535,7 +535,14 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {
   Py_DECREF(py_result);
   PyGILState_Release(gstate);
 
-  unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, num_chunks);
+  // An empty chunk list means this allocation is asleep: its physical chunks
+  // were already unmapped and released by sleep(), but the virtual address is
+  // still held as a placeholder reservation. Skip unmap/release (freeing the
+  // placeholder address happens below).
+  if (num_chunks > 0) {
+    unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes,
+                      num_chunks);
+  }
 #else
   // Non-ROCm path: simple integer handle already extracted; drop temporary
   // Python refs while still holding the GIL, then release it.
@@ -548,11 +555,13 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {
   unmap_and_release(device, size, d_mem, p_memHandle);
 #endif
 
-  // free address and the handle
+  // Free the virtual address. On ROCm this also covers an asleep allocation,
+  // whose placeholder reservation made by sleep() is still held here.
   CUDA_CHECK(cuMemAddressFree(d_mem, size));
 #ifndef USE_ROCM
   free(p_memHandle);
 #else
+  // Only awake allocations have per-chunk handles to free.
   for (auto i = 0; i < num_chunks; ++i) {
     free(p_memHandle[i]);
   }
@@ -672,6 +681,29 @@ static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) {
   unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,
                     num_chunks);
 
+  // On ROCm/Linux, physical VRAM is only reclaimed once the virtual address
+  // range is freed; hipMemUnmap + hipMemRelease alone leave the memory
+  // resident (see ROCm#6021). Free the address to release physical memory,
+  // then immediately re-reserve the SAME address as an empty placeholder so
+  // the regular allocator cannot hand it out while we sleep. wake_up remaps
+  // physical chunks into this placeholder.
+  if (error_code == no_error) {
+    CUDA_CHECK(cuMemAddressFree(d_mem_ptr, recv_size));
+    if (error_code == no_error) {
+      CUdeviceptr reserved = 0;
+      CUDA_CHECK(reserve_rocm_address(&reserved, recv_size, /*alignment=*/0,
+                                      d_mem_ptr));
+      if (error_code == no_error && reserved != d_mem_ptr) {
+        (void)cuMemAddressFree(reserved, recv_size);
+        snprintf(error_msg, sizeof(error_msg),
+                 "failed to re-reserve placeholder address on sleep "
+                 "(requested %#llx, got %#llx)",
+                 (unsigned long long)d_mem_ptr, (unsigned long long)reserved);
+        error_code = CUresult(1);
+      }
+    }
+  }
+
   free(p_memHandle);
   free(chunk_sizes);
 #endif
@@ -736,6 +768,7 @@ static PyObject* python_create_and_map(PyObject* self, PyObject* args) {
     chunk_sizes[i] = PyLong_AsUnsignedLongLong(size_py);
   }
 
+  // Address already reserved as a placeholder by sleep(); just remap chunks.
   create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,
                  num_chunks);
 

csrc/custom_all_reduce_test.cu

@@ -1,361 +0,0 @@
-/**
- * This is a standalone test for custom allreduce.
- * To compile, make sure you have MPI and NCCL installed in your system.
- * export MPI_HOME=XXX
- * nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o
- * custom_all_reduce_test -lnccl -I${MPI_HOME}/include -lmpi
- *
- * Warning: this C++ test is not designed to be very readable and was used
- * during the rapid prototyping process.
- *
- * To run:
- * mpirun --allow-run-as-root -np 8 ./custom_all_reduce_test
- */
-#include <cuda.h>
-#include <curand_kernel.h>
-#include <stdio.h>
-#include <stdlib.h>
-
-#include <limits>
-#include <vector>
-
-#include "cuda_profiler_api.h"
-#include "custom_all_reduce.cuh"
-#include "mpi.h"
-#ifdef USE_ROCM
-  #include <hip/hip_bf16.h>
-typedef __hip_bfloat16 nv_bfloat16;
-  #include "rccl/rccl.h"
-  #include "custom_all_reduce_hip.cuh"
-#else
-  #include "nccl.h"
-  #include "custom_all_reduce.cuh"
-#endif
-
-#define MPICHECK(cmd)                                                  \
-  do {                                                                 \
-    int e = cmd;                                                       \
-    if (e != MPI_SUCCESS) {                                            \
-      printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \
-      exit(EXIT_FAILURE);                                              \
-    }                                                                  \
-  } while (0)
-
-#define NCCLCHECK(cmd)                                              \
-  do {                                                              \
-    ncclResult_t r = cmd;                                           \
-    if (r != ncclSuccess) {                                         \
-      printf("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, \
-             ncclGetErrorString(r));                                \
-      exit(EXIT_FAILURE);                                           \
-    }                                                               \
-  } while (0)
-
-#ifdef USE_ROCM
-__global__ void dummy_kernel() {
-  for (int i = 0; i < 100; i++) {
-    uint64_t start = wall_clock64();
-    uint64_t cycles_elapsed;
-    do {
-      cycles_elapsed = wall_clock64() - start;
-    } while (cycles_elapsed < 100);
-  }
-  for (int i = 0; i < 100; i++) __nanosleep(1000000);  // 100ms
-}
-#else
-__global__ void dummy_kernel() {
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-  for (int i = 0; i < 100; i++) __nanosleep(1000000);  // 100ms
-  #else
-  for (int i = 0; i < 100; i++) {
-    long long int start = clock64();
-    while (clock64() - start < 150000000);  // approximately 98.4ms on P40
-  }
-  #endif
-}
-#endif
-
-template <typename T>
-__global__ void set_data(T* data, int size, int myRank) {
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
-       idx += gridDim.x * blockDim.x) {
-    data[idx] = myRank * 0.11f;
-  }
-}
-
-template <typename T>
-__global__ void convert_data(const T* data1, const T* data2, double* fdata1,
-                             double* fdata2, int size) {
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
-       idx += gridDim.x * blockDim.x) {
-    fdata1[idx] = data1[idx];
-    fdata2[idx] = data2[idx];
-  }
-}
-
-__global__ void init_rand(curandState_t* state, int size, int nRanks) {
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
-       idx += gridDim.x * blockDim.x) {
-    for (int i = 0; i < nRanks; i++) {
-      curand_init(i + 1, idx, 0, &state[idx * nRanks + i]);
-    }
-  }
-}
-
-template <typename T>
-__global__ void gen_data(curandState_t* state, T* data, double* ground_truth,
-                         int myRank, int nRanks, int size) {
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
-       idx += gridDim.x * blockDim.x) {
-    double sum = 0.0;
-    for (int i = 0; i < nRanks; i++) {
-      double val = curand_uniform_double(&state[idx * nRanks + i]) * 4;
-      T hval = val;  // downcast first
-      sum += static_cast<double>(hval);
-      if (i == myRank) data[idx] = hval;
-    }
-    ground_truth[idx] = sum;
-  }
-}
-
-template <typename T>
-void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
-         int data_size, bool performance_test) {
-  T* result;
-  cudaStream_t stream;
-  CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
-  CUDACHECK(cudaMalloc(&result, data_size * sizeof(T)));
-  CUDACHECK(cudaMemset(result, 0, data_size * sizeof(T)));
-
-  cudaIpcMemHandle_t self_data_handle;
-  cudaIpcMemHandle_t data_handles[8];
-  vllm::Signal* buffer;
-  T* self_data_copy;
-  /**
-   * Allocate IPC buffer
-   *
-   * The first section is a temporary buffer for storing intermediate allreduce
-   * results, if a particular algorithm requires it. The second section is for
-   * the input to the allreduce. The actual API takes the input pointer as an
-   * argument (that is, they can and usually should be allocated separately).
-   * But since the input pointers and the temporary buffer all require IPC
-   * registration, they are allocated and registered together in the test for
-   * convenience.
-   */
-#ifdef USE_ROCM
-  CUDACHECK(hipExtMallocWithFlags(
-      (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal),
-      hipDeviceMallocUncached));
-#else
-  CUDACHECK(
-      cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
-#endif
-  CUDACHECK(
-      cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
-  CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T)));
-  CUDACHECK(cudaIpcGetMemHandle(&self_data_handle, buffer));
-
-  MPICHECK(MPI_Allgather(&self_data_handle, sizeof(cudaIpcMemHandle_t),
-                         MPI_BYTE, data_handles, sizeof(cudaIpcMemHandle_t),
-                         MPI_BYTE, MPI_COMM_WORLD));
-
-  void* rank_data;
-  size_t rank_data_sz = 16 * 1024 * 1024;
-  CUDACHECK(cudaMalloc(&rank_data, rank_data_sz));
-  vllm::Signal* ipc_ptrs[8];
-  for (int i = 0; i < nRanks; i++) {
-    if (i == myRank)
-      ipc_ptrs[i] = buffer;
-    else
-      CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i],
-                                     cudaIpcMemLazyEnablePeerAccess));
-  }
-  vllm::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks);
-  auto* self_data =
-      reinterpret_cast<T*>(reinterpret_cast<char*>(buffer) +
-                           sizeof(vllm::Signal) + data_size * sizeof(T));
-  // hack buffer registration
-  {
-    void* data[8];
-    for (int i = 0; i < nRanks; i++) {
-      data[i] =
-          ((char*)ipc_ptrs[i]) + sizeof(vllm::Signal) + data_size * sizeof(T);
-    }
-    fa.register_buffer(data);
-  }
-
-  double* ground_truth;
-  CUDACHECK(cudaMallocHost(&ground_truth, data_size * sizeof(double)));
-  curandState_t* states;
-  CUDACHECK(cudaMalloc(&states, sizeof(curandState_t) * nRanks * data_size));
-  init_rand<<<108, 1024, 0, stream>>>(states, data_size, nRanks);
-  gen_data<T><<<108, 1024, 0, stream>>>(states, self_data, ground_truth, myRank,
-                                        nRanks, data_size);
-  CUDACHECK(cudaMemcpyAsync(self_data_copy, self_data, data_size * sizeof(T),
-                            cudaMemcpyDeviceToDevice, stream));
-  cudaEvent_t start, stop;
-  CUDACHECK(cudaEventCreate(&start));
-  CUDACHECK(cudaEventCreate(&stop));
-
-  ncclDataType_t ncclDtype;
-  if (std::is_same<T, half>::value) {
-    ncclDtype = ncclFloat16;
-  } else if (std::is_same<T, nv_bfloat16>::value) {
-    ncclDtype = ncclBfloat16;
-  } else {
-    ncclDtype = ncclFloat;
-  }
-  double *nccl_result, *my_result;
-  CUDACHECK(cudaMallocHost(&nccl_result, data_size * sizeof(double)));
-  CUDACHECK(cudaMallocHost(&my_result, data_size * sizeof(double)));
-  if (performance_test) {
-    dummy_kernel<<<1, 1, 0, stream>>>();
-    constexpr int warmup_iters = 5;
-    constexpr int num_iters = 100;
-    // warmup
-    for (int i = 0; i < warmup_iters; i++) {
-      NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum,
-                              comm, stream));
-    }
-    CUDACHECK(cudaEventRecord(start, stream));
-    for (int i = 0; i < num_iters; i++) {
-      NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum,
-                              comm, stream));
-    }
-    CUDACHECK(cudaEventRecord(stop, stream));
-    CUDACHECK(cudaStreamSynchronize(stream));
-    float allreduce_ms = 0;
-    cudaEventElapsedTime(&allreduce_ms, start, stop);
-
-    dummy_kernel<<<1, 1, 0, stream>>>();
-    // warm up
-    for (int i = 0; i < warmup_iters; i++) {
-      fa.allreduce<T>(stream, self_data, result, data_size, threads,
-                      block_limit);
-    }
-    CUDACHECK(cudaEventRecord(start, stream));
-    for (int i = 0; i < num_iters; i++) {
-      fa.allreduce<T>(stream, self_data, result, data_size, threads,
-                      block_limit);
-    }
-    CUDACHECK(cudaEventRecord(stop, stream));
-    CUDACHECK(cudaStreamSynchronize(stream));
-
-    float duration_ms = 0;
-    cudaEventElapsedTime(&duration_ms, start, stop);
-    if (myRank == 0)
-      printf(
-          "Rank %d done, nGPUs:%d, sz (kb): %d, %d, %d, my time:%.2fus, nccl "
-          "time:%.2fus\n",
-          myRank, nRanks, data_size * sizeof(T) / 1024, threads, block_limit,
-          duration_ms * 1e3 / num_iters, allreduce_ms * 1e3 / num_iters);
-
-    // And wait for all the queued up work to complete
-    CUDACHECK(cudaStreamSynchronize(stream));
-
-    NCCLCHECK(ncclAllReduce(self_data_copy, self_data, data_size, ncclDtype,
-                            ncclSum, comm, stream));
-
-    convert_data<T><<<108, 1024, 0, stream>>>(self_data, result, nccl_result,
-                                              my_result, data_size);
-    CUDACHECK(cudaStreamSynchronize(stream));
-
-    for (unsigned long j = 0; j < data_size; j++) {
-      auto diff = abs(nccl_result[j] - my_result[j]);
-      if (diff >= 4e-2) {
-        printf("Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
-               myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
-        break;
-      }
-    }
-    long double nccl_diffs = 0.0;
-    long double my_diffs = 0.0;
-    for (int j = 0; j < data_size; j++) {
-      nccl_diffs += abs(nccl_result[j] - ground_truth[j]);
-      my_diffs += abs(my_result[j] - ground_truth[j]);
-    }
-    if (myRank == 0)
-      std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size
-                << " me: " << my_diffs / data_size << std::endl;
-  } else {
-    for (int i = 0; i < 100; i++) {
-      fa.allreduce<T>(stream, self_data, result, data_size, threads,
-                      block_limit);
-      CUDACHECK(cudaStreamSynchronize(stream));
-      NCCLCHECK(ncclAllReduce(self_data, self_data_copy, data_size, ncclDtype,
-                              ncclSum, comm, stream));
-      convert_data<T><<<108, 1024, 0, stream>>>(
-          self_data_copy, result, nccl_result, my_result, data_size);
-      CUDACHECK(cudaStreamSynchronize(stream));
-
-      for (unsigned long j = 0; j < data_size; j++) {
-        auto diff = abs(nccl_result[j] - my_result[j]);
-        if (diff >= 4e-2) {
-          printf(
-              "Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
-              myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
-          break;
-        }
-      }
-    }
-    if (myRank == 0)
-      printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks,
-             data_size * sizeof(T) / 1024, threads, block_limit);
-    // long double nccl_diffs = 0.0;
-    // long double my_diffs = 0.0;
-    // for (int j = 0; j < data_size; j++) {
-    //   nccl_diffs += abs(nccl_result[j] - ground_truth[j]);
-    //   my_diffs += abs(my_result[j] - ground_truth[j]);
-    // }
-    // if (myRank == 0)
-    //   std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size
-    //             << " me: " << my_diffs / data_size << std::endl;
-  }
-
-  CUDACHECK(cudaFree(result));
-  CUDACHECK(cudaFree(self_data_copy));
-  CUDACHECK(cudaFree(rank_data));
-  CUDACHECK(cudaFree(buffer));
-  CUDACHECK(cudaFree(states));
-  CUDACHECK(cudaFreeHost(ground_truth));
-  CUDACHECK(cudaFreeHost(nccl_result));
-  CUDACHECK(cudaFreeHost(my_result));
-  CUDACHECK(cudaStreamDestroy(stream));
-}
-
-int main(int argc, char** argv) {
-  int nRanks, myRank;
-  MPICHECK(MPI_Init(&argc, &argv));
-  MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank));
-  MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks));
-  CUDACHECK(cudaSetDevice(myRank));
-  ncclUniqueId id;
-  ncclComm_t comm;
-  if (myRank == 0) ncclGetUniqueId(&id);
-  MPICHECK(MPI_Bcast(static_cast<void*>(&id), sizeof(id), MPI_BYTE, 0,
-                     MPI_COMM_WORLD));
-  NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank));
-
-  bool performance_test = true;
-  cudaProfilerStart();
-// Uncomment to scan through different block size configs.
-// for (int threads : {256, 512, 1024}) {
-//   for (int block_limit = 16; block_limit < 112; block_limit += 4) {
-//     run<half>(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
-//     performance_test);
-//   }
-// }
-#ifdef USE_ROCM
-  const int block_limit = 16;
-#else
-  const int block_limit = 36;
-#endif
-  // Scan through different sizes to test performance.
-  for (int sz = 512; sz <= (8 << 20); sz *= 2) {
-    run<half>(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test);
-  }
-
-  cudaProfilerStop();
-  MPICHECK(MPI_Finalize());
-  return EXIT_SUCCESS;
-}

csrc/custom_quickreduce.cu

@@ -97,18 +97,28 @@ int64_t qr_max_size() {
                                                   cast_bf2half>;  \
     template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
 
+  // INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the
+  // world_size == 2 kernel for it.
+  #define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \
+    template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, cast_bf2half>;
+
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
+INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
+                                   quickreduce::CodecQ3, false)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
+INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
+                                   quickreduce::CodecQ3, true)
 
 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
+INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false)
 
 #endif  // USE_ROCM

csrc/fs_io.cpp

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#include <Python.h>
+
+#include <unistd.h>
+
+#include <vector>
+
+extern "C" {
+
+static void _batch_lookup(const std::vector<const char*>& paths,
+                          std::vector<int>& exists_flags) {
+  for (size_t i = 0; i < paths.size(); i++) {
+    exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
+  }
+}
+
+/// @brief Check file existence for a batch of paths.
+/// @param paths list[str] – absolute paths to check.
+/// @return list[bool] – True if the corresponding path exists, False otherwise.
+/// @note Releases the GIL for the entire batch. File existence via access(2).
+static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
+  PyObject* path_list;
+  if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) {
+    return nullptr;
+  }
+
+  const Py_ssize_t n = PyList_Size(path_list);
+  std::vector<const char*> paths(n);
+  for (Py_ssize_t i = 0; i < n; i++) {
+    paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr);
+    if (paths[i] == nullptr) {
+      return nullptr;
+    }
+  }
+
+  std::vector<int> exists_flags(n);
+  {
+    Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags);
+    Py_END_ALLOW_THREADS
+  }
+
+  PyObject* result = PyList_New(n);
+  if (result == nullptr) {
+    return nullptr;
+  }
+  for (Py_ssize_t i = 0; i < n; i++) {
+    PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i]));
+  }
+  return result;
+}
+
+static PyMethodDef fs_io_C_methods[] = {
+    {"batch_lookup", batch_lookup, METH_VARARGS,
+     "batch_lookup(paths: list[str]) -> list[bool]\n"
+     "\n"
+     "Check file existence for a batch of paths."},
+    {nullptr, nullptr, 0, nullptr},
+};
+
+static struct PyModuleDef fs_io_C_module = {
+    PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1,
+    fs_io_C_methods,
+};
+
+PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); }
+
+}  // extern "C"

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/cache_kernels.cu

@@ -127,7 +127,12 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
     return reinterpret_cast<BatchFn>(fn_ptr);
   }();
 
-  if (batch_fn != nullptr) {
+  // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 /
+  // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy
+  // fallback below, which is correct on any stream. Real and per-thread-default
+  // streams take the batch fast path.
+  const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
+  if (batch_fn != nullptr && usable_stream) {
     CUmemcpyAttributes attr = {};
     // ANY lets the DMA engine prefetch source bytes out of stream order,
     // which is only safe when no GPU stream is concurrently writing the

csrc/libtorch_stable/cooperative_topk.cu

@@ -0,0 +1,146 @@
+// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
+// See cooperative_topk.cuh for kernel implementation.
+
+#include <cuda_runtime.h>
+
+#include "torch_utils.h"
+
+#ifndef USE_ROCM
+  #include "cooperative_topk.cuh"
+namespace ct = vllm::cooperative;
+namespace hist4096 = vllm::topk_histogram_4096;
+#endif
+
+#ifndef USE_ROCM
+template <uint32_t TopK, uint32_t CS>
+void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
+                                size_t smem, cudaStream_t stream) {
+  auto kernel = []() {
+    if constexpr (CS == 16) {
+      return &ct::cooperative_topk_cs16<TopK>;
+    } else if constexpr (CS == 8) {
+      return &ct::cooperative_topk_cs8<TopK>;
+    } else {
+      static_assert(CS == 4, "unsupported cooperative_topk cluster size");
+      return &ct::cooperative_topk_cs4<TopK>;
+    }
+  }();
+  if constexpr (CS > 8) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed,
+                         1);
+  }
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       smem);
+
+  cudaLaunchConfig_t cfg = {};
+  cfg.gridDim = dim3(params.num_rows, CS);
+  cfg.blockDim = dim3(hist4096::kBlockSize);
+  cfg.dynamicSmemBytes = smem;
+  cfg.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeClusterDimension;
+  attrs[0].val.clusterDim = {1, CS, 1};
+  cfg.numAttrs = 1;
+  cfg.attrs = attrs;
+  cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params);
+  STD_TORCH_CHECK(err == cudaSuccess,
+                  "cooperative_topk launch failed: ", cudaGetErrorString(err));
+}
+
+template <uint32_t TopK>
+void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
+                                  const torch::stable::Tensor& lengths,
+                                  torch::stable::Tensor& output,
+                                  torch::stable::Tensor& workspace,
+                                  int64_t max_seq_len) {
+  (void)max_seq_len;  // Kept for signature parity with persistent_topk.
+  const int64_t num_rows = logits.size(0);
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
+  // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves
+  // headroom)
+  STD_TORCH_CHECK(
+      num_rows <= 32,
+      "cooperative_topk supports <=32 rows; use persistent_topk for "
+      "larger batches");
+
+  STD_TORCH_CHECK(stride % 4 == 0,
+                  "cooperative_topk: stride must be multiple of 4 for TMA "
+                  "alignment, got stride (max_model_len)=",
+                  stride);
+
+  STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
+  STD_TORCH_CHECK(
+      workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
+      "workspace must be uint8");
+
+  ct::CooperativeTopKParams<TopK> params;
+  params.input = logits.const_data_ptr<float>();
+  params.output = output.mutable_data_ptr<int32_t>();
+  params.lengths = lengths.const_data_ptr<int32_t>();
+  params.num_rows = static_cast<uint32_t>(num_rows);
+  params.stride = stride;
+  params.tie_ws =
+      reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  STD_TORCH_CHECK(
+      workspace.size(0) >=
+          static_cast<int64_t>(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)),
+      "workspace too small");
+
+  const bool supports_cluster16 = get_device_prop()->major >= 10;
+  if (num_rows <= 4 && supports_cluster16) {
+    launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
+  } else if (num_rows <= 8) {
+    launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
+  } else {
+    launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
+  }
+}
+#endif  // USE_ROCM
+
+void cooperative_topk(const torch::stable::Tensor& logits,
+                      const torch::stable::Tensor& lengths,
+                      torch::stable::Tensor& output,
+                      torch::stable::Tensor& workspace, int64_t k,
+                      int64_t max_seq_len) {
+#ifndef USE_ROCM
+  STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "Only float32 supported");
+  STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "lengths must be int32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "output must be int32");
+  STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+                  "lengths must be 1D or 2D");
+  STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
+  const int64_t num_rows = logits.size(0);
+  STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+                  "output size mismatch");
+  STD_TORCH_CHECK(
+      k == 512 || k == 1024 || k == 2048,
+      "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k);
+
+  if (k == 512) {
+    launch_cooperative_topk_impl<512>(logits, lengths, output, workspace,
+                                      max_seq_len);
+  } else if (k == 1024) {
+    launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  } else {
+    launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  }
+#else
+  STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm");
+#endif
+}

csrc/libtorch_stable/cooperative_topk.cuh

@@ -0,0 +1,593 @@
+/*
+ * Cooperative TopK kernel for DSA Indexer
+ */
+
+#ifndef COOPERATIVE_TOPK_CUH_
+#define COOPERATIVE_TOPK_CUH_
+
+#include <cooperative_groups.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cuda/ptx>
+#include <algorithm>
+#include <cstdint>
+
+#include "topk_histogram_4096.cuh"
+
+namespace vllm {
+namespace cooperative {
+
+namespace hist4096 = topk_histogram_4096;
+
+constexpr uint32_t kHistBits = 10;
+constexpr uint32_t kHistBins = 1 << kHistBits;
+constexpr uint32_t kMaxTopK = 2048;
+
+constexpr uint32_t kElemPerStage = 16;
+constexpr uint32_t kSizePerStage =
+    kElemPerStage * hist4096::kBlockSize;  // 16384
+
+// CS=4 two-pass path uses two TMA stages as a double buffer.
+constexpr uint32_t kStreamingStagesCS4 = 2;
+// CS=8/16 fused paths keep all loaded TMA stages resident in smem.
+constexpr uint32_t kFusedStagesCS8 = 2;
+constexpr uint32_t kFusedStagesCS16 = 2;
+
+// CS=4 single-pass path
+constexpr uint32_t kMaxSinglePassStages = 3;
+constexpr uint32_t kMaxSinglePassPerBlock =
+    kMaxSinglePassStages * kSizePerStage;  // 49152
+
+template <uint32_t TopK = 1024>
+struct CooperativeTopKParams {
+  const float* __restrict__ input;
+  int32_t* __restrict__ output;
+  const int32_t* __restrict__ lengths;
+  hist4096::Tie* __restrict__ tie_ws;  // per-row tie workspace, see
+                                       // kTieWsPerRow
+  uint32_t num_rows, stride;
+};
+
+// ============================================================================
+// Cooperative helpers
+// ============================================================================
+
+// only CS adjacent lanes participate (sub-warp reduce), in opposite to
+// warp_reduce_sum_full
+template <uint32_t N>
+__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) {
+#pragma unroll
+  for (uint32_t m = N >> 1; m > 0; m >>= 1)
+    v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32);
+  return v;
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+__device__ __forceinline__ uint32_t extract_coarse_bin(float x) {
+  return hist4096::extract_coarse_bin_N<kHistBits>(x);
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) {
+  cuda::ptx::mbarrier_init(a, n);
+}
+__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) {
+  while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
+                                              cuda::ptx::scope_cta, a, p));
+}
+__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a,
+                                                          uint32_t t) {
+  cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
+                                       cuda::ptx::scope_cta,
+                                       cuda::ptx::space_shared, a, t);
+}
+__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n,
+                                         uint64_t* m) {
+  cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d,
+                           s, n, m);
+}
+
+// ============================================================================
+// DSMEM histogram reduce
+// ============================================================================
+
+template <uint32_t CS>
+__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) {
+  static_assert(kHistBins <= hist4096::kBlockSize);
+  auto cluster = cooperative_groups::this_cluster();
+  cluster.sync();
+  const auto tx = threadIdx.x;
+  const auto rank = blockIdx.y;
+  constexpr auto kLocal = kHistBins / CS;
+  const auto off = kLocal * rank;
+  if (tx < kHistBins) {
+    const auto addr = &histogram[off + tx / CS];
+    const auto src = cluster.map_shared_rank(addr, tx % CS);
+    *src = warp_reduce_sum_subN<CS>(*src);
+  }
+  cluster.sync();
+}
+
+// ============================================================================
+// Find threshold from reduced histogram
+// ============================================================================
+
+// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened
+// before calling this, so warp_sum writes are visible across warps.
+// The first internal __syncthreads() is still needed for the warp_sum exchange.
+template <uint32_t TopK>
+__device__ __forceinline__ void find_threshold(uint32_t* histogram,
+                                               uint32_t* warp_sum,
+                                               uint32_t* counter_gt,
+                                               uint32_t* counter_eq,
+                                               hist4096::MatchBin* match) {
+  const auto tx = threadIdx.x;
+  const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize;
+  const auto value = tx < kHistBins ? histogram[tx] : 0;
+  const auto winc = hist4096::warp_inclusive_sum(li, value);
+  if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc;
+  __syncthreads();
+  const auto tmp = warp_sum[li];
+  const auto total = hist4096::warp_reduce_sum_full(tmp);
+  auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc;
+  const auto above = total - pfx;
+  if (tx < kHistBins && above < TopK && above + value >= TopK) {
+    *counter_gt = *counter_eq = 0;
+    *match = {.bin = tx, .above_count = above, .equal_count = value};
+  }
+  __syncthreads();
+}
+
+// Streams data through shared memory in chunks, processing each chunk before
+// loading the next overwrites each buffer after processing it (the epilogue
+// prefetch loads the next chunk into the same slot)
+template <typename SmemType, uint32_t kStages, uint32_t kBinBits,
+          bool kIsScatter>
+__device__ void tma_stream_pass(const float* scores, uint32_t length,
+                                uint32_t thr_bin, int32_t* indices,
+                                uint32_t* phases, SmemType* smem) {
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+  const auto ni =
+      (length + kSizePerStage - 1) / kSizePerStage;  // total stages needed
+  const auto la =
+      (length + 3u) & ~3u;  // length rounded up to float4 (TMA alignment)
+  const auto pass =
+      kIsScatter ? 1 : 0;  // barrier dim: [0] for histogram, [1] for scatter
+
+  // Prologue: issue initial TMA loads - prefill the pipeline
+  if (tx == 0) {
+#pragma unroll
+    for (uint32_t i = 0; i < kStages; i++) {
+      if (i >= ni) {
+        break;
+      }
+      const auto o = i * kSizePerStage;
+      const auto sz = min(kSizePerStage, la - o) * sizeof(float);
+      tma_load(smem->score_buffer[i], scores + o, sz,
+               &smem->barrier[pass][i]);  // cp.async.bulk is non-blocking
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz);
+    }
+  }
+
+  // Main loop: process stages
+  for (uint32_t it = 0; it < ni; it++) {
+    const auto b = it % kStages;  // which buffer slot (0 or 1)
+    const auto o = it * kSizePerStage;
+    const auto sz = min(kSizePerStage, length - o);
+
+    if (lane == 0) {
+      mbarrier_wait(&smem->barrier[pass][b],
+                    phases[b] & 1);  // wait for the data
+    }
+    phases[b]++;  // advances the phase for next time this slot is reused
+    __syncwarp();
+
+#pragma unroll
+    for (uint32_t i = 0; i < kElemPerStage; i++) {
+      const auto li = tx + i * hist4096::kBlockSize;
+      if (li >= sz) {
+        break;
+      }
+      const auto sc = smem->score_buffer[b][li];
+      const auto bn = hist4096::extract_coarse_bin_N<kBinBits>(sc);
+      if constexpr (kIsScatter) {  // compile-time branch
+        // Scatter pass: place above-threshold and collect ties
+        const auto gi = o + li;
+        if (bn > thr_bin) {
+          indices[atomicAdd(&smem->counter_gt, 1)] = gi;
+        } else if (bn == thr_bin) {
+          const auto p = atomicAdd(&smem->counter_eq, 1);
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gi, sc};
+          }
+        }
+      } else {
+        // Histogram pass: just count
+        atomicAdd(&smem->histogram[bn], 1);
+      }
+    }
+    __syncthreads();  // ensures all threads finished processing their buffer
+                      // before next TMA load
+
+    // Epilogue: issue next TMA load
+    if (tx == 0 && it + kStages < ni) {
+      const auto no = (it + kStages) * kSizePerStage;
+      const auto nsz = min(kSizePerStage, la - no) * sizeof(float);
+      tma_load(smem->score_buffer[b], scores + no, nsz,
+               &smem->barrier[pass][b]);
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz);
+    }
+  }
+}
+
+// ============================================================================
+// Fused path: single TMA pass, rescan smem for scatter
+// ============================================================================
+
+// Fused shared memory layout for cluster cooperative paths.
+// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass
+// (CS=4).
+template <uint32_t kStages, uint32_t kPasses = 1>
+struct SmemFused {
+  uint64_t barrier[kPasses][kStages];
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  alignas(128) hist4096::MatchBin match;
+  uint32_t warp_sum[hist4096::kNumWarps];
+  union {
+    uint32_t histogram[kHistBins];
+    hist4096::Tie tie_buffer[kMaxTopK];
+  };
+  alignas(128) float score_buffer[kStages][kSizePerStage];
+};
+
+using Smem8 = SmemFused<kFusedStagesCS8>;
+using Smem16 = SmemFused<kFusedStagesCS16>;
+using Smem4 = SmemFused<kStreamingStagesCS4, 2>;
+using SmemSinglePass = SmemFused<kMaxSinglePassStages>;
+
+// Cluster-cooperative large path.
+// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan
+// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram
+// then scatter).
+template <uint32_t TopK, uint32_t CS, typename SmemType, bool kFused>
+__device__ void large_topk(const float* __restrict__ row_input,
+                           int32_t* __restrict__ row_output, uint32_t seq_len,
+                           uint32_t* phases, hist4096::Tie* tie_ws) {
+  const auto rank = blockIdx.y;  // this block's position in cluster
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+
+  extern __shared__ uint8_t smem_raw[];
+  auto* smem = reinterpret_cast<SmemType*>(smem_raw);
+  int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw + sizeof(SmemType));
+
+  // Partition row across cluster ranks
+  constexpr uint32_t kAlign = 4;
+  const auto units =
+      (seq_len + kAlign - 1) / kAlign;  // float4-aligned element count
+  const auto base = units / CS, extra = units % CS;  // elements per block
+  const auto lu = base + (rank < extra ? 1u : 0u);   // remainder blocks
+  const auto ou =
+      rank * base + min(rank, extra);  // this block's count (load-balanced)
+  const auto my_start = ou * kAlign;   // global start offset
+  const auto my_len = min(my_start + lu * kAlign, seq_len) -
+                      my_start;  // actual length of this block
+  const auto num_iters =
+      (my_len + kSizePerStage - 1) / kSizePerStage;  // TMA stages needed
+  const auto len_aligned = (my_len + 3u) & ~3u;
+
+  if constexpr (kFused) {
+    // Fused init + TMA prologue
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;  // all threads zero histogram
+    }
+    if (tx == 0) {  // thread 0 issues TMA - then all threads continue working
+                    // until mbarrier sync
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+      for (uint32_t i = 0; i < num_iters; i++) {
+        const auto off = i * kSizePerStage;
+        const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float);
+        tma_load(smem->score_buffer[i], row_input + my_start + off, sz,
+                 &smem->barrier[0][i]);  // cp.async.bulk of size kSizePerStage
+                                         // × sizeof(float)
+        mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz);
+      }
+    }
+    __syncthreads();
+
+    // Histogram build. ILP unroll-by-2, no inter-stage sync
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+      if (lane == 0) {
+        mbarrier_wait(&smem->barrier[0][iter],
+                      phases[iter] & 1);  // wait for TMA
+      }
+      phases[iter]++;
+      __syncwarp();
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i += 2) {
+        const auto li0 = tx + i * hist4096::kBlockSize;
+        const auto li1 = tx + (i + 1) * hist4096::kBlockSize;
+        if (li0 >= sz) {
+          break;
+        }
+        const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]);
+        if (li1 < sz) {
+          const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]);
+          atomicAdd(&smem->histogram[b0], 1);
+          atomicAdd(&smem->histogram[b1], 1);
+        } else {
+          atomicAdd(&smem->histogram[b0], 1);
+        }
+      }
+    }
+  } else {
+    // Twopass: init then stream histogram pass
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;
+    }
+    if (tx == 0) {
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+    }
+    __syncthreads();
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, false>(
+        row_input + my_start, my_len, 0, nullptr, phases, smem);
+  }
+
+  // DSMEM all-reduce + find threshold
+  dsmem_hist_reduce<CS>(
+      smem->histogram);  // each block histogram is summed across all CS blocks
+  find_threshold<TopK>(smem->histogram, smem->warp_sum, &smem->counter_gt,
+                       &smem->counter_eq, &smem->match);
+
+  const auto thr = smem->match.bin;
+
+  if constexpr (kFused) {
+    // Fused scatter: rescan score_buffer (still in smem)
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i++) {
+        const auto li = tx + i * hist4096::kBlockSize;
+        if (li >= sz) {
+          break;
+        }
+        const auto score = smem->score_buffer[iter][li];  // still in smem
+        const auto bin = extract_coarse_bin(score);
+        const auto gidx = off + li;
+        if (bin > thr) {
+          s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx;  // above -> s_topk
+        } else if (bin == thr) {
+          const auto p = atomicAdd(&smem->counter_eq,
+                                   1);  // equal -> ties (later refinement)
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gidx, score};
+          }
+        }
+      }
+    }
+    __syncthreads();
+  } else {
+    // Twopass scatter: re-stream data via TMA
+    uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0};
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, true>(
+        row_input + my_start, my_len, thr, s_topk, scatter_phases, smem);
+  }
+
+  // Output collection via DSMEM prefix sum
+  constexpr uint32_t kAboveBits = 16;
+  constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
+  static_assert(kAboveMask >= TopK);
+  static_assert(kAboveMask >= kMaxSinglePassPerBlock,
+                "kAboveBits must cover max per-block element count");
+
+  const uint32_t la = smem->counter_gt;
+  const uint32_t le_full = smem->counter_eq;
+  const uint32_t le =
+      min(le_full, hist4096::kMaxTies);  // written smem tie_buffer entries
+
+  __shared__ uint32_t s_local_counts[CS];
+  __shared__ uint32_t s_prefix_packed;
+  __shared__ uint32_t s_total_above, s_total_equal;
+
+  auto cluster = cooperative_groups::this_cluster();
+  if (tx < CS) {
+    // Pack written tie counts into 32-bit: (equal << 16) | above.
+    // `le_full` may exceed the per-block tie buffer cap; using it here creates
+    // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots.
+    const uint32_t packed = (le << kAboveBits) | la;
+    const auto dst = cluster.map_shared_rank(s_local_counts, tx);
+    dst[rank] = packed;  // write my count to every block's s_local_counts[rank]
+  }
+  cluster.sync();
+
+  // Thread 0 computes serial prefix sum
+  if (tx == 0) {
+    uint32_t prefix = 0, ta = 0, te = 0;
+    for (uint32_t i = 0; i < CS; i++) {
+      if (i == rank) {
+        s_prefix_packed = prefix;  // my prefix
+      }
+      ta += s_local_counts[i] & kAboveMask;   // total above
+      te += s_local_counts[i] >> kAboveBits;  // total equal
+      prefix += s_local_counts[i];
+    }
+    s_total_above = ta;
+    s_total_equal = te;
+  }
+  __syncthreads();
+
+  const uint32_t prefix_above = s_prefix_packed & kAboveMask;
+  const uint32_t prefix_equal = s_prefix_packed >> kAboveBits;
+
+  // Write to global output
+  for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) {
+    // indices are placed contiguously starting at prefix_above
+    row_output[prefix_above + i] =
+        s_topk[i] + my_start;  // my_start: block-local -> row-global index
+  }
+  for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) {
+    const auto t = smem->tie_buffer[i];
+    uint32_t p = s_total_above + prefix_equal + i;
+    if (p < TopK) {
+      row_output[p] = t.idx + my_start;
+    }
+    uint32_t tp = prefix_equal + i;
+    if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) {
+      tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score};
+    }
+  }
+
+  // Tie refinement
+  cooperative_groups::this_cluster().sync();
+  if (rank != 0) {  // only rank 0 does tie refinement
+    return;
+  }
+  if (s_total_above + s_total_equal <= TopK) {  // no ties to refine
+    return;
+  }
+
+  // Tie-breaking uses FP32 (4-round radix sort)
+  if constexpr (TopK <= hist4096::kBlockSize) {
+    // copy ties from tie_ws back to smem, then refine
+    const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies);
+    // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie)
+    for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) {
+      smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score};
+    }
+    __syncthreads();
+    hist4096::tie_handle<TopK>(smem->tie_buffer, num_ties, s_total_above,
+                               row_output, smem);
+  } else {
+    // TopK=2048: process directly from tie_ws (GMEM)
+    const uint32_t num_ties = min(s_total_equal, static_cast<uint32_t>(TopK));
+    hist4096::tie_handle_large<TopK>(tie_ws, num_ties, s_total_above,
+                                     row_output, smem);
+  }
+}
+
+// ============================================================================
+// Adapted from https://github.com/sgl-project/sglang/pull/23600
+// sgl-project/sglang
+// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/)
+// ============================================================================
+
+template <uint32_t TopK, uint32_t CS>
+__device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
+  const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
+  const auto sl = params.lengths[row];
+  int32_t* out = params.output + row * TopK;
+  const float* in = params.input + row * params.stride;
+
+  // Trivial: seq_len <= TopK
+  if (sl <= static_cast<int32_t>(TopK)) {
+    if (rank == 0) {
+      for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) {
+        out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1;
+      }
+    }
+    return;
+  }
+
+  // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF
+  if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) {
+    if (rank == 0) {
+      extern __shared__ uint8_t sr[];
+      hist4096::histogram_4096_topk<TopK, 12>(
+          in, out, sl, sr);  // 4096-bin (12-bit) histogram
+    }
+    return;
+  }
+
+  // Large path: init mbarriers + state, then dispatch fused or twopass
+  const uint32_t per_block =
+      (params.stride + CS - 1) / CS;  // how many elements per block
+  constexpr uint32_t kFusedMax = ((CS == 16)  ? kFusedStagesCS16
+                                  : (CS == 8) ? kFusedStagesCS8
+                                              : kMaxSinglePassStages) *
+                                 kSizePerStage;
+  const bool use_singlepass =
+      per_block <=
+      kFusedMax;  // single pass or TMA streaming: histogram+scatter
+
+  // Select smem type and stage count at compile time based on CS
+  constexpr uint32_t kFusedStages = (CS == 16)  ? kFusedStagesCS16
+                                    : (CS == 8) ? kFusedStagesCS8
+                                                : kMaxSinglePassStages;
+  using FusedSmem = SmemFused<kFusedStages>;
+
+  extern __shared__ uint8_t sr[];
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow;
+
+  if (use_singlepass) {
+    auto* smem = reinterpret_cast<FusedSmem*>(sr);
+    const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage;
+    if (tx < sp_stages) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 1 barrier per TMA stage -
+                         // signal when async copies complete
+    }
+    __syncthreads();
+    uint32_t phases[kFusedStages] =
+        {};  // tracks the parity for mbarrier wait/arrive protocol
+    large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
+  } else {
+    // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
+    auto* smem = reinterpret_cast<Smem4*>(sr);
+    if (tx < 2 * kStreamingStagesCS4) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 2×2=4 barriers (2 passes × 2 stages)
+    }
+    __syncthreads();
+    uint32_t hp[kStreamingStagesCS4] = {0,
+                                        0};  // histogram+scatter pass counters
+    large_topk<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws);
+  }
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 4, 1)
+        cooperative_topk_cs4(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 4>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 8, 1)
+        cooperative_topk_cs8(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 8>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 16, 1)
+        cooperative_topk_cs16(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 16>(params);
+}
+
+constexpr size_t kSmemSize4_base = sizeof(Smem4);
+constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
+constexpr size_t kSmemSize4 =
+    (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
+    sizeof(int32_t) * 2048 + 128;
+constexpr size_t kSmemSize8 =
+    sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
+
+}  // namespace cooperative
+
+}  // namespace vllm
+
+#endif  // COOPERATIVE_TOPK_CUH_