Bump actions/setup-python from 6.1.0 to 6.3.0

Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6.1.0 to 6.3.0.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/83679a892e2d95755f2dac6acb0bfd1e9ac5d548...ece7cb06caefa5fff74198d8649806c4678c61a1)

---
updated-dependencies:
- dependency-name: actions/setup-python
  dependency-version: 6.3.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>

.buildkite/ci_config_intel.yaml

@@ -2,17 +2,16 @@ name: vllm_intel_ci
 job_dirs:
   - ".buildkite/intel_jobs"
 run_all_patterns:
+  - ".buildkite/ci_config_intel.yaml"
   - "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: []
@@ -136,7 +140,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 3
+  parallelism: 4
 
 - label: "Arm CPU Test"
   depends_on: []

.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml

@@ -21,6 +21,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -38,6 +42,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -55,6 +63,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"

.buildkite/image_build/image_build_arm64.sh

@@ -21,12 +21,13 @@ else
   exit 0
 fi
 
-# build (Grace/GH200 is the arm64 GPU target; sm_90)
+# build for arm64 GPU targets: Grace/GH200 (sm_90) and DGX Spark/GB10
+# (sm_121, family-covered by 12.0 under CUDA 13)
 docker build --file docker/Dockerfile \
   --platform linux/arm64 \
   --build-arg max_jobs=16 \
   --build-arg nvcc_threads=4 \
-  --build-arg torch_cuda_arch_list="9.0" \
+  --build-arg torch_cuda_arch_list="9.0 12.0" \
   --build-arg USE_SCCACHE=1 \
   --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
   --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \

.buildkite/intel_jobs/basic_correctness.yaml

@@ -5,6 +5,10 @@ steps:
 - label: XPU Sleep Mode
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -19,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/engine_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: Engine (1 GPU)
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/expert_parallelism_intel.yaml

@@ -1,11 +1,15 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: EPLB Algorithm
   key: eplb-algorithm
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/kernels_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: vLLM IR Tests
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/lora_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: LoRA Runtime + Utils
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -34,6 +38,10 @@ steps:
 - label: LoRA Fused/MoE Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -54,6 +62,10 @@ steps:
 - label: LoRA Punica Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -74,6 +86,10 @@ steps:
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -94,6 +110,10 @@ steps:
 - label: LoRA Models
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -108,15 +128,19 @@ 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
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/misc_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: V1 Core + KV + Metrics
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -31,6 +35,10 @@ steps:
 - label: V1 Sample + Logits
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -71,6 +79,10 @@ steps:
 - label: XPU CPU Offload
   timeout_in_minutes: 60
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -95,6 +107,10 @@ steps:
   key: regression
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -126,6 +142,10 @@ steps:
   timeout_in_minutes: 30
   num_devices: 2
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -157,6 +177,10 @@ steps:
   key: async-engine-inputs-utils-worker
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/model_runner_v2_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: Model Runner V2 Core Tests (Intel)
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -30,6 +34,10 @@ steps:
 - label: Model Runner V2 Examples (Intel)
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:

.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

@@ -1,11 +1,15 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
   key: multi-modal-models-standard-1-qwen2
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -27,6 +31,10 @@ steps:
   key: multi-modal-models-standard-2-qwen3-gemma
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -47,6 +55,10 @@ steps:
   key: multi-modal-models-standard-3-llava-qwen2-vl
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -68,6 +80,10 @@ steps:
   key: multi-modal-models-standard-4-other-whisper
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -88,6 +104,10 @@ steps:
   key: multi-modal-processor
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -105,7 +125,5 @@ steps:
       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/test-intel.yaml

@@ -19,6 +19,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -49,6 +53,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -74,6 +82,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -93,6 +105,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"

.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,24 @@ 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
+if git -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+  vllm_standalone_merge_base="$(
+    git -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+  )"
+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}}"
@@ -525,6 +557,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

@@ -34,7 +34,7 @@ case "${test_suite}" in
     pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
     pytest -v -s v1/structured_output
     pytest -v -s v1/test_serial_utils.py
-    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
     pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
     ;;
   server)

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

@@ -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/install-kv-connectors.sh

@@ -4,6 +4,11 @@
 
 set -euo pipefail
 
+if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
+    uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+    exit 0
+fi
+
 REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
 
 uv pip install --system -r "${REQUIREMENTS_FILE}"

.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/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,8 +225,8 @@ 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/"

.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
@@ -74,6 +76,20 @@ steps:
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py
 
 - label: Kernels Attention DiffKV Test (H100)
   key: kernels-attention-diffkv-test-h100
@@ -128,6 +144,22 @@ steps:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 5
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      source_file_dependencies:
+      - csrc/quantization/cutlass_w8a8/moe/
+      - csrc/moe/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/device_communicators/
+      - vllm/envs.py
+      - vllm/config
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 - label: Kernels Mamba Test
   key: kernels-mamba-test
@@ -242,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
@@ -259,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 (2xB200)
+  key: lm-eval-small-models-2xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -64,8 +65,20 @@ 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 (2xL4)
+  key: lm-eval-small-models-tp
+  timeout_in_minutes: 10
+  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 +89,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 +106,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
@@ -109,6 +122,7 @@ steps:
       - image-build-amd
       commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
       - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
 
 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
@@ -135,6 +149,49 @@ steps:
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
+- label: LM Eval Humming (A100 - TEMPORARY)
+  key: lm-eval-humming-a100
+  timeout_in_minutes: 30
+  device: a100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+
+- label: LM Eval Humming (H100 - TEMPORARY)
+  key: lm-eval-humming-h100
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+
+- label: LM Eval Humming (B200 - TEMPORARY)
+  key: lm-eval-humming-b200
+  timeout_in_minutes: 30
+  device: b200-k8s
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
 
 - label: LM Eval TurboQuant KV Cache
   key: lm-eval-turboquant-kv-cache
@@ -148,8 +205,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
@@ -162,8 +219,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
@@ -176,6 +233,22 @@ 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-b200.txt
 
+- label: GPQA Eval (GPT-OSS) (DGX Spark)
+  key: gpqa-eval-gpt-oss-spark
+  timeout_in_minutes: 120
+  device: dgx-spark
+  optional: true
+  num_devices: 1
+  depends_on:
+  - arm64-image-build
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/evals/gpt_oss/
+  commands:
+    - 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-spark.txt
+
 - label: MRCR Eval Small Models
   device: h200_35gb
   timeout_in_minutes: 30

.buildkite/test_areas/misc.yaml

@@ -105,6 +105,12 @@ steps:
     # Integration test for streaming correctness (requires special branch).
     - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
 
 - label: V1 Others (CPU)
   key: v1-others-cpu
@@ -218,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
@@ -244,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
@@ -256,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
@@ -338,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/models_basic.yaml

@@ -45,6 +45,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

@@ -15,6 +15,10 @@ steps:
     - 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
@@ -32,6 +36,21 @@ steps:
   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

.buildkite/test_areas/models_multimodal.yaml

@@ -109,6 +109,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

.buildkite/test_areas/pytorch.yaml

@@ -107,6 +107,12 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
 
 - label: PyTorch Fullgraph Smoke Test
   key: pytorch-fullgraph-smoke-test
@@ -189,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

@@ -12,6 +12,20 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py
 
 - label: Spec Decode Eagle Nightly B200
   key: spec-decode-eagle-nightly-b200
@@ -80,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,8 +48,8 @@ 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/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+    - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
       with:
         python-version: "3.12"
     # Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its

.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

AGENTS.md

@@ -135,6 +135,12 @@ Do not modify code in these areas without first reading and following the
 linked guide. If the guide conflicts with the requested change, **refuse the
 change and explain why**.
 
+Security reviewers should start with [`SECURITY.md`](SECURITY.md),
+[`docs/usage/security.md`](docs/usage/security.md), and
+[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
+for the project security policy, threat model, deployment assumptions, and
+vulnerability process.
+
 - **Editing these instructions**:
   [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
   — Rules for modifying AGENTS.md or any domain-specific guide it references.

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()
 
 #
@@ -382,6 +407,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "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")
 
@@ -498,6 +541,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 +894,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")
@@ -1049,6 +1100,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     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)

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/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)
@@ -423,6 +427,7 @@ 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_fused_moe.cpp"
         ${VLLM_EXT_SRC})
 endif()
 

cmake/external_projects/deepgemm.cmake

@@ -8,43 +8,73 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR})
   set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
 endif()
 
+# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download).
+# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW
+# disallows one-argument Populate(dep) after Declare; MakeAvailable would run
+# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load).
 if(DEEPGEMM_SRC_DIR)
-  FetchContent_Declare(
-    deepgemm
-    SOURCE_DIR ${DEEPGEMM_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # cmake_path(ABSOLUTE_PATH <var> ...) reads the path from <var>; NORMALIZE is a
+  # flag (no trailing path argument). Resolve relative paths against vLLM root.
+  set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _deepgemm_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}")
+  if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}")
+    message(FATAL_ERROR
+      "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'")
+  endif()
+  set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
+  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
-  # This ref should be kept in sync with tools/install_deepgemm.sh
-  FetchContent_Declare(
-    deepgemm
-    GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
-    GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # Keep in sync with tools/install_deepgemm.sh
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+
+  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _deepgemm_fc_root)
+    set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src")
+  set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build")
+  set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild")
+
+  if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp")
+    set(deepgemm_SOURCE_DIR "${_deepgemm_src}")
+    set(deepgemm_BINARY_DIR "${_deepgemm_bin}")
+  else()
+    FetchContent_Populate(
+      deepgemm
+      SUBBUILD_DIR "${_deepgemm_sub}"
+      SOURCE_DIR "${_deepgemm_src}"
+      BINARY_DIR "${_deepgemm_bin}"
+      GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
+      GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
+      GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+  message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
 endif()
 
-# Use FetchContent_Populate (not MakeAvailable) to avoid processing
-# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
-FetchContent_GetProperties(deepgemm)
-if(NOT deepgemm_POPULATED)
-  FetchContent_Populate(deepgemm)
-endif()
-message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
-
-# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
+# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream),
+# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x
+# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants.
 set(DEEPGEMM_SUPPORT_ARCHS)
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
   list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
 endif()
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
-elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a")
+  endif()
 endif()
 
 cuda_archs_loose_intersection(DEEPGEMM_ARCHS

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/qutlass.cmake

@@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
   set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
 endif()
 
+# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
+# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
 if(QUTLASS_SRC_DIR)
-  FetchContent_Declare(
-    qutlass
-    SOURCE_DIR ${QUTLASS_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _qutlass_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
+  if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
+    message(FATAL_ERROR
+      "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
+  endif()
+  set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
+  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
-  FetchContent_Declare(
-    qutlass
-    GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
-    GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
-endif()
+  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
 
-FetchContent_Populate(qutlass)
+  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _qutlass_fc_root)
+    set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
+  set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
+  set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
+
+  if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
+    set(qutlass_SOURCE_DIR "${_qutlass_src}")
+    set(qutlass_BINARY_DIR "${_qutlass_bin}")
+  else()
+    FetchContent_Populate(
+      qutlass
+      SUBBUILD_DIR "${_qutlass_sub}"
+      SOURCE_DIR "${_qutlass_src}"
+      BINARY_DIR "${_qutlass_bin}"
+      GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
+      GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+endif()
 
 if(NOT qutlass_SOURCE_DIR)
   message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")

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 b3964b1d8b95d8e8447435668ab169a2700bab65
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

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;
@@ -887,12 +887,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 +1423,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_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

@@ -497,6 +497,26 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
     reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
   };
 
+  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 +535,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 +962,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

@@ -4,6 +4,9 @@
 #ifdef CPU_CAPABILITY_AMXBF16
   #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
 #endif
+#if defined(__riscv_v)
+  #include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
+#endif
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
 
 #define VLLM_DISPATCH_CASE_16B_TYPES(...)                 \
@@ -152,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
@@ -319,6 +322,8 @@ void cpu_gemm_wna16(
       return ISA::AMX;
     } else if (isa_hint == "vec") {
       return ISA::VEC;
+    } else if (isa_hint == "rvv") {
+      return ISA::RVV;
     } else {
       TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
     }
@@ -397,6 +402,40 @@ void cpu_gemm_wna16(
             pack_factor);
         return;
       }
+    } else if (isa == ISA::RVV) {
+      using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
+      if (has_zp) {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      }
+      if (use_desc_act) {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      } else {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      }
     }
   });
 }

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_rvv.hpp

@@ -0,0 +1,228 @@
+#ifndef CPU_MICRO_GEMM_RVV_HPP
+#define CPU_MICRO_GEMM_RVV_HPP
+
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+#if defined(__riscv_v)
+
+namespace cpu_micro_gemm {
+namespace {
+
+constexpr int32_t RVV_MGEMM_N8 = 8;
+constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
+
+template <typename scalar_t>
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
+
+template <>
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
+}
+
+template <>
+FORCE_INLINE fixed_fp32x8_t
+load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
+  #if defined(__riscv_zvfh)
+  fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
+      reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
+  return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
+  #else
+  alignas(32) float values[RVV_MGEMM_N8];
+  for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
+    values[i] = static_cast<float>(ptr[i]);
+  }
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
+  #endif
+}
+
+template <>
+FORCE_INLINE fixed_fp32x8_t
+load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
+  #if defined(__riscv_zvfbfmin)
+  fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
+      reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
+  fixed_bf16x8_t vec =
+      RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
+  return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
+  #else
+  fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
+      reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
+  auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
+  auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
+  return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
+  #endif
+}
+
+// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
+// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
+template <int32_t M, typename scalar_t>
+FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
+                                             const scalar_t* __restrict__ b_ptr,
+                                             float* __restrict__ c_ptr,
+                                             const int64_t lda,
+                                             const int64_t ldc, const int32_t k,
+                                             const bool accum_c) {
+  static_assert(0 < M && M <= 8);
+
+  #define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
+  #define RVV_IF_M(i) if constexpr (M > (i))
+
+  #define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
+  RVV_ROWS_APPLY(RVV_DECL_A)
+  #undef RVV_DECL_A
+
+  #define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
+  RVV_ROWS_APPLY(RVV_DECL_ACC)
+  #undef RVV_DECL_ACC
+
+  #define RVV_INIT_ACC(i)                                                  \
+    RVV_IF_M(i) {                                                          \
+      if (accum_c) {                                                       \
+        acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc,    \
+                                                     RVV_MGEMM_N8);        \
+      } else {                                                             \
+        acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
+      }                                                                    \
+    }
+  RVV_ROWS_APPLY(RVV_INIT_ACC)
+  #undef RVV_INIT_ACC
+
+  int32_t k_idx = 0;
+  for (; k_idx + 3 < k; k_idx += 4) {
+  #define RVV_FMA_ROW(i, K_OFFSET)                                     \
+    RVV_IF_M(i) {                                                      \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(                  \
+          acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
+          RVV_MGEMM_N8);                                               \
+    }
+
+  #define RVV_STEP_K(K_OFFSET)                                      \
+    {                                                               \
+      fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>(              \
+          b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
+      RVV_FMA_ROW(0, K_OFFSET)                                      \
+      RVV_FMA_ROW(1, K_OFFSET)                                      \
+      RVV_FMA_ROW(2, K_OFFSET)                                      \
+      RVV_FMA_ROW(3, K_OFFSET)                                      \
+      RVV_FMA_ROW(4, K_OFFSET)                                      \
+      RVV_FMA_ROW(5, K_OFFSET)                                      \
+      RVV_FMA_ROW(6, K_OFFSET)                                      \
+      RVV_FMA_ROW(7, K_OFFSET)                                      \
+    }
+
+    RVV_STEP_K(0)
+    RVV_STEP_K(1)
+    RVV_STEP_K(2)
+    RVV_STEP_K(3)
+  #undef RVV_STEP_K
+  #undef RVV_FMA_ROW
+  }
+
+  for (; k_idx < k; ++k_idx) {
+    fixed_fp32x8_t b =
+        load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
+  #define RVV_TAIL_ROW(i)                                                \
+    RVV_IF_M(i) {                                                        \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(                    \
+          acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
+    }
+    RVV_ROWS_APPLY(RVV_TAIL_ROW)
+  #undef RVV_TAIL_ROW
+  }
+
+  #define RVV_STORE_ROW(i)                                           \
+    RVV_IF_M(i) {                                                    \
+      RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
+                                          RVV_MGEMM_N8);             \
+    }
+  RVV_ROWS_APPLY(RVV_STORE_ROW)
+  #undef RVV_STORE_ROW
+
+  #undef RVV_ROWS_APPLY
+  #undef RVV_IF_M
+}
+
+template <int32_t M, typename scalar_t>
+FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+  static_assert(0 < M && M <= 8);
+  scalar_t* __restrict__ curr_b_0 = b_ptr;
+  scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
+
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
+                                c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
+                                accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
+                                ldc, k, accum_c);
+}
+
+class TileGemmRVV {
+ public:
+  template <typename scalar_t>
+  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    switch (m) {
+      case 1:
+        gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 2:
+        gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 3:
+        gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 4:
+        gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 5:
+        gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 6:
+        gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 7:
+        gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 8:
+        gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
+        break;
+    }
+  }
+};
+
+}  // namespace
+
+template <typename scalar_t>
+class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
+ public:
+  static constexpr int32_t MaxMSize = 8;
+  static constexpr int32_t NSize = 32;
+
+ public:
+  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
+  }
+
+  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_EQ(output_size % 16, 0);
+    for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
+      const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
+      scalar_t* __restrict__ curr_packed_weight =
+          packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
+      for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
+        *curr_packed_weight = *curr_weight;
+
+        curr_packed_weight += 16;
+        ++curr_weight;
+      }
+    }
+  }
+};
+
+}  // namespace cpu_micro_gemm
+
+#endif  // defined(__riscv_v)
+
+#endif  // CPU_MICRO_GEMM_RVV_HPP

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,
@@ -538,7 +548,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 +559,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 +609,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,19 +2,24 @@
 #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 };
+enum class ISA { AMX, VEC, RVV, NEON };
 
 inline ISA get_isa(const std::string& isa) {
   if (isa == "amx") {
     return ISA::AMX;
   } else if (isa == "vec") {
     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/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/mla/sm100_cutlass_mla_kernel.cu

@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
   StrideQ stride_Q_pe = cute::make_tuple(
       static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
 
+  // Read the token and page strides from the cache tensor instead of assuming
+  // packed pages, so strided views (e.g. per-layer views into a cross-layer
+  // block-major cache) are addressed correctly.
   StrideK stride_C = cute::make_tuple(
-      static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
   StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
   StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
   StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));

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
@@ -549,7 +554,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
     const int head_dim,                        // dimension of each head
     const int quant_block_size,                // quantization block size
     const int cache_block_size,                // cache block size
-    const int cache_stride,  // stride for each token in kv_cache
+    const int64_t cache_block_stride,  // stride for each block in kv_cache
 
     const bool use_ue8m0  // use ue8m0 scale format
 ) {
@@ -590,16 +595,15 @@ __global__ void indexer_k_quant_and_cache_kernel(
     scale = exp2f(ceilf(log2f(scale)));
   }
 
-  const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
-                             block_offset * head_dim + head_dim_idx;
+  const int64_t dst_offset =
+      block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
   for (int i = 0; i < VEC_SIZE; i++) {
     kv_cache[dst_offset + i] =
         fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
   }
   if (threadIdx.x == 0) {
     const int64_t dst_scale_idx =
-        block_idx * cache_block_size * cache_stride +
-        cache_block_size * head_dim +
+        block_idx * cache_block_stride + cache_block_size * head_dim +
         (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
     reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
   }
@@ -1452,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
           reinterpret_cast<KV_T*>(k.data_ptr()),                              \
           reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),                    \
           slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
-          cache_block_size, cache_stride, use_ue8m0);
+          cache_block_size, cache_block_stride, use_ue8m0);
 
 void indexer_k_quant_and_cache(
     torch::stable::Tensor& k,         // [num_tokens, head_dim]
@@ -1463,7 +1467,7 @@ void indexer_k_quant_and_cache(
   int num_tokens = k.size(0);
   int head_dim = k.size(1);
   int cache_block_size = kv_cache.size(1);
-  int cache_stride = kv_cache.size(2);
+  int64_t cache_block_stride = kv_cache.stride(0);
   bool use_ue8m0 = scale_fmt == "ue8m0";
 
   STD_TORCH_CHECK(k.device() == kv_cache.device(),

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_

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -67,6 +67,13 @@
   #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
+// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
+#ifndef USE_ROCM
+  #include <cuda_fp8.h>
+#else
+  #include <hip/hip_fp8.h>
+#endif
+
 #ifndef FINAL_MASK
   #ifdef USE_ROCM
     #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,6 +82,19 @@
   #endif
 #endif
 
+#ifdef USE_ROCM
+// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
+// fused kernel).
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
+  #if defined(HIP_FP8_TYPE_OCP)
+  __hip_fp8_e4m3 fp8_val(val);
+  #else
+  __hip_fp8_e4m3_fnuz fp8_val(val);
+  #endif
+  return reinterpret_cast<uint8_t&>(fp8_val);
+}
+#endif
+
 namespace vllm {
 namespace minimax_m3_fused_ops {
 
@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
   *reinterpret_cast<uint2*>(dst) = v;
 }
 
+// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
+// dtypes use the scaled-convert path with identity scale.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __device__ __forceinline__ void storeCacheElems(
     cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
   }
 }
 
+// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
+// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
+// (RMSNorm outputs are O(1) and the score path only needs relative block
+// ordering).
+__device__ __forceinline__ void storeElemsFp8(
+    uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  constexpr float kFp8Max = 448.0f;
+#ifndef USE_ROCM
+  __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
+    vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
+    vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
+    out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
+  }
+  *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
+#else
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane; i++) {
+    float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
+    dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
+  }
+#endif
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
 //     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
 //     IQ: niq  only if kIsSparse        (norm+RoPE)
 //     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
+// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          bool kIsSparse, bool kInsertKV>
+          typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
-    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
-    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
+    out_idx_t* __restrict__ index_q_out,  // [N, niq*128]; scalar_t or e4m3 byte
     scalar_t const* __restrict__ q_norm_w,
     scalar_t const* __restrict__ k_norm_w,
     scalar_t const* __restrict__ iq_norm_w,
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int64_t const* __restrict__ positions,       // [N] i64
     int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
     int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
-    cache_t* __restrict__ kv_cache,      // [nb,2,bs,nkv,128] or nullptr
-    scalar_t* __restrict__ index_cache,  // [nb*bs, 128] or nullptr
+    cache_t* __restrict__ kv_cache,       // [nb,2,bs,nkv,128] or nullptr
+    out_idx_t* __restrict__ index_cache,  // [nb*bs, 128]; scalar_t or e4m3 byte
     float const eps, int const rotary_dim, int const num_tokens, int const nq,
     int const nkv, int const niq, int const block_size,
     // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
                   slot * kHeadDim;
     } else if (isIQ && index_q_out != nullptr) {
-      store_ptr = index_q_out +
-                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
-                  (slot - iq_begin) * kHeadDim;
+      // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
+      if constexpr (!kFp8Idx) {
+        store_ptr = index_q_out +
+                    static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                    (slot - iq_begin) * kHeadDim;
+      }
     }
 
     // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
       normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
                             cos_ptr, /*apply_norm=*/norm_w != nullptr);
-      storeElems<scalar_t>(store_ptr + dim_base, elems);
+      if constexpr (kFp8Idx) {
+        // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
+        if (isIQ && index_q_out != nullptr) {
+          storeElemsFp8(index_q_out +
+                            static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                            (slot - iq_begin) * kHeadDim + dim_base,
+                        elems);
+        } else {
+          storeElems<scalar_t>(store_ptr + dim_base, elems);
+        }
+      } else {
+        storeElems<scalar_t>(store_ptr + dim_base, elems);
+      }
     }
 
     // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
                              : (isIK ? index_slot_mapping[tokenIdx] : -1);
       if (sm >= 0) {  // skip padded / unscheduled tokens
         if (isIK) {
-          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
-          storeElems<scalar_t>(dst, elems);
+          if constexpr (kFp8Idx) {
+            storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
+          } else {
+            storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
+          }
         } else if (isK || isV) {
           // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
           // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
-void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
-                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
-                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
-                          scalar_t const* cos_sin_cache,
-                          int64_t const* positions, int64_t const* slot_mapping,
-                          int64_t const* index_slot_mapping, cache_t* kv_cache,
-                          scalar_t* index_cache, float const eps,
-                          int const rotary_dim, int const num_tokens,
-                          int const nq, int const nkv, int const niq,
-                          int const block_size, int64_t const kv_s_block,
-                          int64_t const kv_s_kv, int64_t const kv_s_token,
-                          int64_t const kv_s_head, bool const has_index,
-                          bool const insert_kv, cudaStream_t stream) {
+void launchFusedMiniMaxM3(
+    scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
+    scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
+    scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
+    int64_t const* positions, int64_t const* slot_mapping,
+    int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head, bool const has_index, bool const insert_kv,
+    bool const fp8_idx, cudaStream_t stream) {
+  // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
+  // void* pointers per instantiation in the LAUNCH macro.
   // Slot count must match the kernel's compile-time gating.
   int const v_slots = insert_kv ? nkv : 0;
   int const idx_slots = has_index ? niq + 1 : 0;
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  #define LAUNCH(IS_SPARSE, INSERT)                                           \
-    cudaLaunchKernelEx(                                                       \
-        &config,                                                              \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,       \
-                                              IS_SPARSE, INSERT>,             \
-        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,    \
-        cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
-        index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size,   \
-        kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                                \
+    cudaLaunchKernelEx(                                                        \
+        &config,                                                               \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
+                                              IS_SPARSE, INSERT, FP8>,         \
+        qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
+        iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping,          \
+        index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache),   \
+        eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,     \
+        kv_s_kv, kv_s_token, kv_s_head)
 #else
   // ROCm: standard kernel launch syntax (no PDL/stream serialization).
   // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT)                                          \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,          \
-                                          IS_SPARSE, INSERT>                 \
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                              \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T,   \
+                                          IS_SPARSE, INSERT, FP8>            \
         <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
-            ik_norm_w, cos_sin_cache, positions, slot_mapping,               \
-            index_slot_mapping, kv_cache, index_cache, eps, rotary_dim,      \
+            qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,     \
+            k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,        \
+            slot_mapping, index_slot_mapping, kv_cache,                      \
+            reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim,          \
             num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
             kv_s_token, kv_s_head)
   // clang-format on
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
 
   if (has_index) {
     if (insert_kv) {
-      LAUNCH(true, true);  // sparse serving
+      if (fp8_idx) {
+        LAUNCH(true, true, true, uint8_t);  // sparse serving, fp8 index outputs
+      } else {
+        LAUNCH(true, true, false, scalar_t);  // sparse serving, bf16
+      }
     } else {
-      LAUNCH(true, false);  // sparse profiling
+      if (fp8_idx) {
+        LAUNCH(true, false, true, uint8_t);  // sparse profiling, fp8 index_q
+      } else {
+        LAUNCH(true, false, false, scalar_t);  // sparse profiling, bf16
+      }
     }
   } else {
     // Dense layer: never has an index branch and never inserts here (the
     // generic Attention layer owns the KV insert).
-    LAUNCH(false, false);
+    LAUNCH(false, false, false, scalar_t);
   }
 #undef LAUNCH
 }
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
   vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
       reinterpret_cast<st*>(qkv.data_ptr()),                                   \
       q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr,  \
-      index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
-                              : nullptr,                                       \
+      index_q_out.has_value()                                                  \
+          ? reinterpret_cast<void*>(index_q_out->data_ptr())                   \
+          : nullptr,                                                           \
       reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
       reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
       has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
                 : nullptr,                                                     \
       insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
       (insert_kv && has_index)                                                 \
-          ? reinterpret_cast<st*>(index_cache->data_ptr())                     \
+          ? reinterpret_cast<void*>(index_cache->data_ptr())                   \
           : nullptr,                                                           \
       static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
       nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
-      kv_s_head, has_index, insert_kv, stream)
+      kv_s_head, has_index, insert_kv, fp8_idx, stream)
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                  torch::headeronly::ScalarType::Long &&
              index_slot_mapping->numel() == slot_mapping->numel()),
         "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
     if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
       STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
                       "auto kv_cache dtype must match qkv");
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
           kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
           "fp8 kv_cache must use uint8 storage");
     }
-    STD_TORCH_CHECK(index_cache.has_value() &&
-                        index_cache->scalar_type() == qkv.scalar_type(),
-                    "insert mode requires matching index_cache");
+    // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
+    STD_TORCH_CHECK(
+        index_cache.has_value() &&
+            (index_cache->scalar_type() == qkv.scalar_type() ||
+             index_cache->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
     STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
                     "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
                     "head_dim (stride(4)==1)");
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
         "index_q_out requires the index branch (num_index_heads > 0)");
     STD_TORCH_CHECK(
         index_q_out->is_cuda() && index_q_out->is_contiguous() &&
-            index_q_out->scalar_type() == qkv.scalar_type(),
-        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+            (index_q_out->scalar_type() == qkv.scalar_type() ||
+             index_q_out->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
     STD_TORCH_CHECK(index_q_out->numel() ==
                         static_cast<int64_t>(num_tokens) * niq * kHeadDim,
                     "index_q_out must have num_tokens * num_index_heads * 128 "
                     "elements");
   }
 
+  // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
+  // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
+  auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
+  bool const fp8_idx =
+      (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
+      (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
+  if (fp8_idx) {
+    STD_TORCH_CHECK(
+        !index_cache.has_value() || index_cache->scalar_type() == kFp8,
+        "fp8 index path: index_cache must be fp8 e4m3");
+    STD_TORCH_CHECK(
+        !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
+        "fp8 index path: index_q_out must be fp8 e4m3");
+  }
+
   const torch::stable::accelerator::DeviceGuard device_guard(
       qkv.get_device_index());
   auto stream = get_current_cuda_stream(qkv.get_device_index());

csrc/libtorch_stable/layernorm_kernels.cu

@@ -20,20 +20,26 @@ __global__ void rms_norm_kernel(
     const int64_t input_stride_d4,        // input.stride(-4)
     const int64_t input_shape_d2,         // input.size(-2)
     const int64_t input_shape_d3,         // input.size(-3)
-    const scalar_t* __restrict__ weight,  // [hidden_size], null if !HasWeight
+    const scalar_t* __restrict__ weight,  // [hidden_size] or
+                                          // [num_groups, hidden_size];
+                                          // null if !HasWeight
+    const int64_t weight_stride,          // 0 or weight.stride(0)
     const float epsilon, const int num_tokens, const int hidden_size) {
   __shared__ float s_variance;
   float variance = 0.0f;
   const scalar_t* input_row;
+  const scalar_t* weight_row;
   if constexpr (NUM_DIMS == 2) {
     // 2D for layernorm normal case [batch_size, hidden]
     input_row = input + blockIdx.x * input_stride_d2;
+    weight_row = weight + blockIdx.x * weight_stride;
   } else if constexpr (NUM_DIMS == 3) {
     // 3D for q/k norm [batch_size, num_heads, head_size]
     int batch_idx = blockIdx.x / input_shape_d2;
     int head_idx = blockIdx.x % input_shape_d2;
     input_row =
         input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
+    weight_row = weight + batch_idx * weight_stride;
   } else if constexpr (NUM_DIMS == 4) {
     // 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
     int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2);
@@ -42,6 +48,7 @@ __global__ void rms_norm_kernel(
     int head_idx = remaining % input_shape_d2;
     input_row = input + batch_idx * input_stride_d4 +
                 seq_idx * input_stride_d3 + head_idx * input_stride_d2;
+    weight_row = weight + batch_idx * weight_stride;
   }
 
   auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
@@ -69,7 +76,7 @@ __global__ void rms_norm_kernel(
 
   scalar_t* out_row = out + blockIdx.x * hidden_size;
   auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
-  auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
+  auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight_row);
   auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row);
   for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) {
     vec_n_t<scalar_t, VEC_SIZE> dst;
@@ -211,15 +218,24 @@ fused_add_rms_norm_kernel(
 
 void rms_norm(torch::stable::Tensor& out,    // [..., hidden_size]
               torch::stable::Tensor& input,  // [..., hidden_size]
-              std::optional<torch::stable::Tensor> weight,  // [hidden_size]
-              double epsilon) {
+              std::optional<torch::stable::Tensor> weight, double epsilon) {
   STD_TORCH_CHECK(out.is_contiguous());
   if (input.stride(-1) != 1) {
     input = torch::stable::contiguous(input);
   }
   STD_TORCH_CHECK(input.stride(-1) == 1);
+  int64_t weight_stride = 0;
   if (weight.has_value()) {
     STD_TORCH_CHECK(weight->is_contiguous());
+    if (weight->dim() == 1) {
+      STD_TORCH_CHECK(weight->size(0) == input.size(-1));
+    } else if (weight->dim() == 2) {
+      STD_TORCH_CHECK(weight->size(0) == input.size(0));
+      STD_TORCH_CHECK(weight->size(-1) == input.size(-1));
+      weight_stride = weight->stride(0);
+    } else {
+      STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D");
+    }
   }
 
   int hidden_size = input.size(-1);
@@ -256,16 +272,16 @@ void rms_norm(torch::stable::Tensor& out,    // [..., hidden_size]
                       out.mutable_data_ptr<scalar_t>(),
                       input.const_data_ptr<scalar_t>(), input_stride_d2,
                       input_stride_d3, input_stride_d4, input_shape_d2,
-                      input_shape_d3, weight_ptr, epsilon, num_tokens,
-                      hidden_size);
+                      input_shape_d3, weight_ptr, weight_stride, epsilon,
+                      num_tokens, hidden_size);
             } else {
               vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, false>
                   <<<grid, block, 0, stream>>>(
                       out.mutable_data_ptr<scalar_t>(),
                       input.const_data_ptr<scalar_t>(), input_stride_d2,
                       input_stride_d3, input_stride_d4, input_shape_d2,
-                      input_shape_d3, weight_ptr, epsilon, num_tokens,
-                      hidden_size);
+                      input_shape_d3, weight_ptr, /*weight_stride=*/0, epsilon,
+                      num_tokens, hidden_size);
             }
           });
         });

csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu

@@ -286,3 +286,52 @@ template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
 
 template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
     __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

@@ -41,6 +41,7 @@ inline int getSMVersion() {
 static constexpr int DEFAULT_NUM_EXPERTS = 256;
 static constexpr int KIMI_K2_NUM_EXPERTS = 384;
 static constexpr int DEFAULT_HIDDEN_DIM = 7168;
+static constexpr int GLM_5_HIDDEN_DIM = 6144;
 
 template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
@@ -121,14 +122,21 @@ void dsv3_router_gemm(
 
   STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
                   "mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
-                  "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
-                  ", but got hidden_dim=", hidden_dim);
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM,
+      "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
+      " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim);
   STD_TORCH_CHECK(
       num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
       "Expected num_experts=", DEFAULT_NUM_EXPERTS,
       " or num_experts=", KIMI_K2_NUM_EXPERTS,
       ", but got num_experts=", num_experts);
+  // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim.
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS,
+      "hidden_dim=", GLM_5_HIDDEN_DIM,
+      " only supports num_experts=", DEFAULT_NUM_EXPERTS,
+      ", but got num_experts=", num_experts);
   STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                   "currently num_tokens must be less than or equal to 16 for "
                   "router_gemm");
@@ -148,35 +156,49 @@ void dsv3_router_gemm(
 
   const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
 
+  __nv_bfloat16 const* a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  __nv_bfloat16 const* b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
   if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
+                                                          a_ptr, b_ptr, stream);
     }
   } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    __nv_bfloat16* out_ptr =
+        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
+                                                         a_ptr, b_ptr, stream);
     }
   }
 }

csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu

@@ -286,3 +286,52 @@ template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
 
 template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
     float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -11,6 +11,7 @@
 #include "../../cuda_compat.h"
 #include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/dispatch_utils.h"
+#include "libtorch_stable/quantization/vectorization.cuh"
 #include "libtorch_stable/torch_utils.h"
 
 #define CEILDIV(x, y) (((x) + (y) - 1) / (y))
@@ -349,19 +350,102 @@ __global__ void count_and_sort_expert_tokens_kernel(
       max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
 }
 
+// Reduce the topk expert outputs per token (summed in fp32). The output is
+// dense [num_tokens, d]; the input is addressed by its strides so non-
+// contiguous inputs work without a copy. A 16B-vectorized path is used when
+// the hidden dim is contiguous (innermost stride 1) and aligned; otherwise a
+// scalar kernel reads via arbitrary strides. topk is a compile-time constant
+// for common values and runtime otherwise.
+
+// Elements per 16-byte vector (8 for bf16/fp16, 4 for fp32).
+template <typename scalar_t>
+constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t);
+
 template <typename scalar_t, int TOPK>
-__global__ void moe_sum_kernel(
-    scalar_t* __restrict__ out,          // [..., d]
-    const scalar_t* __restrict__ input,  // [..., topk, d]
-    const int d) {
-  const int64_t token_idx = blockIdx.x;
-  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
-    scalar_t x = 0.0;
+__global__ void moe_sum_vec_kernel(
+    scalar_t* __restrict__ out,          // [num_tokens, d], contiguous
+    const scalar_t* __restrict__ input,  // [num_tokens, topk, d], d contiguous
+    const int64_t num_tokens, const int d, const int64_t stride_token,
+    const int64_t stride_topk) {
+  using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>;  // 16-byte pack
+  constexpr int VEC = MOE_SUM_VEC<scalar_t>;
+  const int64_t n_vec = d / VEC;
+  const int64_t total = num_tokens * n_vec;
+  for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total;
+       i += (int64_t)gridDim.x * blockDim.x) {
+    const int64_t token = i / n_vec;
+    const int64_t v = i % n_vec;
+    const scalar_t* in_tok = input + token * stride_token + v * VEC;
+
+    float acc[VEC];
+#pragma unroll
+    for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
+
 #pragma unroll
     for (int k = 0; k < TOPK; ++k) {
-      x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
+      vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
+#pragma unroll
+      for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
     }
-    out[token_idx * d + idx] = x;
+
+    vec_t outp;
+#pragma unroll
+    for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]);
+    *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp;
+  }
+}
+
+// Runtime-topk variant of the above.
+template <typename scalar_t>
+__global__ void moe_sum_vec_dynamic_kernel(
+    scalar_t* __restrict__ out,          // [num_tokens, d], contiguous
+    const scalar_t* __restrict__ input,  // [num_tokens, topk, d], d contiguous
+    const int64_t num_tokens, const int d, const int topk,
+    const int64_t stride_token, const int64_t stride_topk) {
+  using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>;
+  constexpr int VEC = MOE_SUM_VEC<scalar_t>;
+  const int64_t n_vec = d / VEC;
+  const int64_t total = num_tokens * n_vec;
+  for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total;
+       i += (int64_t)gridDim.x * blockDim.x) {
+    const int64_t token = i / n_vec;
+    const int64_t v = i % n_vec;
+    const scalar_t* in_tok = input + token * stride_token + v * VEC;
+
+    float acc[VEC];
+#pragma unroll
+    for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
+
+    for (int k = 0; k < topk; ++k) {
+      vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
+#pragma unroll
+      for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
+    }
+
+    vec_t outp;
+#pragma unroll
+    for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]);
+    *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp;
+  }
+}
+
+// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims
+// (including a non-contiguous hidden stride) via per-element strided reads.
+template <typename scalar_t>
+__global__ void moe_sum_scalar_kernel(
+    scalar_t* __restrict__ out,          // [num_tokens, d], contiguous
+    const scalar_t* __restrict__ input,  // [num_tokens, topk, d]
+    const int d, const int topk, const int64_t stride_token,
+    const int64_t stride_topk, const int64_t stride_hidden) {
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* in_tok = input + token_idx * stride_token;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    float x = 0.f;
+    for (int k = 0; k < topk; ++k) {
+      x += static_cast<float>(
+          VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden]));
+    }
+    out[token_idx * d + idx] = static_cast<scalar_t>(x);
   }
 }
 
@@ -626,52 +710,82 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
 void moe_sum(torch::stable::Tensor& input,   // [num_tokens, topk, hidden_size]
              torch::stable::Tensor& output)  // [num_tokens, hidden_size]
 {
+  // Output is dense and written in place, so it must be contiguous. The input
+  // is read by its strides (no copy); only the hidden dim needs to be
+  // contiguous to take the vectorized path.
+  STD_TORCH_CHECK(output.is_contiguous(),
+                  "moe_sum expects a contiguous output");
+
   const int hidden_size = input.size(-1);
-  const auto num_tokens = output.numel() / hidden_size;
+  const int64_t num_tokens = output.numel() / hidden_size;
   const int topk = input.size(1);
+  const int64_t stride_token = input.stride(0);
+  const int64_t stride_topk = input.stride(1);
+  const int64_t stride_hidden = input.stride(2);
 
-  dim3 grid(num_tokens);
-  dim3 block(std::min(hidden_size, 1024));
   const torch::stable::accelerator::DeviceGuard device_guard(
       output.get_device_index());
   const cudaStream_t stream =
       get_current_cuda_stream(output.get_device_index());
 
-  switch (topk) {
-    case 2:
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-          input.scalar_type(), "moe_sum_kernel", [&] {
-            vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
-                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
-                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
-                hidden_size);
-          });
-      break;
+#define LAUNCH_MOE_SUM_VEC(TOPK)                \
+  vllm::moe::moe_sum_vec_kernel<scalar_t, TOPK> \
+      <<<grid, dim3(block), 0, stream>>>(       \
+          out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk)
 
-    case 3:
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-          input.scalar_type(), "moe_sum_kernel", [&] {
-            vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
-                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
-                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
-                hidden_size);
-          });
-      break;
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] {
+    constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
+    constexpr int WIDTH = VEC * sizeof(scalar_t);  // 16 bytes
+    auto* out_ptr = reinterpret_cast<scalar_t*>(output.mutable_data_ptr());
+    auto* in_ptr = reinterpret_cast<const scalar_t*>(input.const_data_ptr());
 
-    case 4:
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-          input.scalar_type(), "moe_sum_kernel", [&] {
-            vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
-                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
-                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
-                hidden_size);
-          });
-      break;
-
-    default:
-      torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
-      break;
-  }
+    // Vectorize along hidden only when it is contiguous (innermost stride 1),
+    // a whole number of vectors, and every row offset stays 16B-aligned.
+    const bool can_vec = (stride_hidden == 1) && (hidden_size % VEC == 0) &&
+                         (stride_token % VEC == 0) &&
+                         (stride_topk % VEC == 0) &&
+                         (reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) &&
+                         (reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0);
+    if (can_vec) {
+      const int64_t n_vec = hidden_size / VEC;
+      const int64_t total = num_tokens * n_vec;
+      const int block = 256;
+      const dim3 grid(std::min<int64_t>((total + block - 1) / block, 65535));
+      switch (topk) {
+        case 1:
+          LAUNCH_MOE_SUM_VEC(1);
+          break;
+        case 2:
+          LAUNCH_MOE_SUM_VEC(2);
+          break;
+        case 4:
+          LAUNCH_MOE_SUM_VEC(4);
+          break;
+        case 6:
+          LAUNCH_MOE_SUM_VEC(6);
+          break;
+        case 8:
+          LAUNCH_MOE_SUM_VEC(8);
+          break;
+        case 9:
+          LAUNCH_MOE_SUM_VEC(9);
+          break;
+        default:
+          vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t>
+              <<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens,
+                                                 hidden_size, topk,
+                                                 stride_token, stride_topk);
+          break;
+      }
+    } else {
+      dim3 grid(num_tokens);
+      dim3 block(std::min(hidden_size, 1024));
+      vllm::moe::moe_sum_scalar_kernel<scalar_t><<<grid, block, 0, stream>>>(
+          out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk,
+          stride_hidden);
+    }
+  });
+#undef LAUNCH_MOE_SUM_VEC
 }
 
 void moe_lora_align_block_size(

csrc/libtorch_stable/ops.h

@@ -343,6 +343,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
                      torch::stable::Tensor& workspace, int64_t k,
                      int64_t max_seq_len);
 
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+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);
+#endif
+
 void selective_scan_fwd(
     const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
     const torch::stable::Tensor& A, const torch::stable::Tensor& B,