wip

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
[MM][CG] Gemma3 Encoder CUDA Graph (#43591 )
2026-06-29 18:15:10 +00:00 · 2026-06-29 04:52:00 +00:00 · 2026-06-29 03:39:04 +00:00 · 2026-06-29 03:04:05 +00:00 · 2026-06-28 19:43:58 -07:00 · 2026-06-29 08:11:17 +08:00
794 changed files with 40665 additions and 14033 deletions
@@ -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"
@@ -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
@@ -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: []
@@ -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
@@ -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}"
@@ -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
@@ -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 .
@@ -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

@@ -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 \
@@ -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
@@ -112,33 +112,6 @@ steps:
 #                                                                                                                                       #
 #########################################################################################################################################

-#-----------------------------------------------------  mi250 · basic_correctness  -----------------------------------------------------#
-
- label: Distributed Model Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/model_loader/sharded_state_loader.py
-  - vllm/model_executor/models/
-  - vllm/model_executor/layers/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - tests/basic_correctness/
-  - tests/model_executor/model_loader/test_sharded_state_loader.py
-  - tests/models/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
-
 #----------------------------------------------------------  mi250 · compile  ----------------------------------------------------------#

 - label: PyTorch Compilation Unit Tests # TBD
@@ -179,72 +152,8 @@ steps:
  commands:
  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"

- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/compile/test_wrapper.py
-  - tests/entrypoints/llm/test_collective_rpc.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
-
 #--------------------------------------------------------  mi250 · distributed  --------------------------------------------------------#

- label: Distributed Comm Ops # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed
-  - tests/distributed
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s distributed/test_comm_ops.py
-  - pytest -v -s distributed/test_shm_broadcast.py
-  - pytest -v -s distributed/test_shm_buffer.py
-  - pytest -v -s distributed/test_shm_storage.py
-
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/distributed/
-  - tests/v1/shutdown
-  - tests/v1/worker/test_worker_memory_snapshot.py
-  - vllm/platforms/rocm.py
-  commands:
-  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
-  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
-
 - label: Pipeline + Context Parallelism (4 GPUs) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -354,54 +263,6 @@ steps:
  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
  - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model

-#----------------------------------------------------------  mi250 · plugins  ----------------------------------------------------------#
-
- label: Plugin Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/plugins/
-  - tests/plugins/
-  - vllm/platforms/rocm.py
-  commands:
-  # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform
-  - pip install -e ./plugins/vllm_add_dummy_platform
-  - pytest -v -s plugins_tests/test_platform_plugins.py
-  - pip uninstall vllm_add_dummy_platform -y
-  # END: platform plugin tests
-  # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin
-  - pip install -e ./plugins/prithvi_io_processor_plugin
-  - pytest -v -s plugins_tests/test_io_processor_plugins.py
-  - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
-  - pip uninstall prithvi_io_processor_plugin -y
-  # END: `io_processor` plugins test
-  # BEGIN: `bge_m3_sparse io_processor` test
-  - pip install -e ./plugins/bge_m3_sparse_plugin
-  - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
-  - pip uninstall bge_m3_sparse_plugin -y
-  # END: `bge_m3_sparse io_processor` test
-  # BEGIN: `colbert_query io_processor` test
-  - pip install -e ./plugins/colbert_query_plugin
-  - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
-  - pip uninstall colbert_query_plugin -y
-  # END: `colbert_query io_processor` test
-  # BEGIN: `stat_logger` plugins test
-  - pip install -e ./plugins/vllm_add_dummy_stat_logger
-  - pytest -v -s plugins_tests/test_stats_logger_plugins.py
-  - pip uninstall dummy_stat_logger -y
-  # END: `stat_logger` plugins test
-  # BEGIN: other tests
-  - pytest -v -s plugins_tests/test_scheduler_plugins.py
-  - pip install -e ./plugins/vllm_add_dummy_model
-  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
-  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
-  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
-
 #------------------------------------------------------------  mi250 · v1  -------------------------------------------------------------#

 - label: Batch Invariance (H100-MI250) # TBD
@@ -529,61 +390,6 @@ steps:
  commands:
  - pytest -v -s v1/attention

- label: V1 Sample + Logits # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/v1/sample
-  - tests/v1/logits_processors
-  - tests/v1/test_oracle.py
-  - tests/v1/test_request.py
-  - tests/v1/test_outputs.py
-  commands:
-  - pytest -v -s v1/sample
-  - pytest -v -s v1/logits_processors
-  - pytest -v -s v1/test_oracle.py
-  - pytest -v -s v1/test_request.py
-  - pytest -v -s v1/test_outputs.py
-
- label: Distributed DP Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  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
-  commands:
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - 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
-
- label: V1 e2e (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/v1/e2e
-  commands:
-    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
-
 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#

 - label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD
@@ -617,7 +423,7 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - 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 parser
  - pytest -v -s transformers_utils
@@ -632,7 +438,7 @@ steps:
 #-----------------------------------------------------  mi300 · basic_correctness  -----------------------------------------------------#

 - label: Basic Correctness # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
@@ -646,11 +452,11 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s basic_correctness/test_mem.py
-  - pytest -v -s basic_correctness/test_basic_correctness.py
+  - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py
  - pytest -v -s basic_correctness/test_cpu_offload.py

 - label: Distributed Model Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 65
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -668,7 +474,7 @@ steps:
  - tests/model_executor/model_loader/test_sharded_state_loader.py
  - tests/models/
  commands:
-  - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+  - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
@@ -682,6 +488,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -794,7 +601,7 @@ steps:

 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#

- label: Platform Tests (CUDA) # TBD
+- label: Platform Tests # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
@@ -825,10 +632,27 @@ steps:

 #--------------------------------------------------------  mi300 · distributed  --------------------------------------------------------#

+- label: Distributed Comm Ops # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed
+  - tests/distributed
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s distributed/test_comm_ops.py
+  - pytest -v -s distributed/test_shm_broadcast.py
+  - pytest -v -s distributed/test_shm_buffer.py
+  - pytest -v -s distributed/test_shm_storage.py
+
 - label: EPLB Algorithm # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/distributed/eplb
@@ -853,8 +677,8 @@ steps:
  - pytest -v -s distributed/test_eplb_execute.py
  - pytest -v -s distributed/test_eplb_spec_decode.py

- label: Distributed Tests (2xH100-2xMI250) # TBD
-  timeout_in_minutes: 180
+- label: Distributed Tests (2xH100-2xMI300) # TBD
+  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -865,13 +689,19 @@ steps:
  - vllm/model_executor/layers/fused_moe/
  - vllm/v1/attention/backends/
  - vllm/v1/attention/selector.py
+  - tests/v1/distributed/test_dbo.py
  - tests/distributed/test_context_parallel.py
  - examples/features/data_parallel/data_parallel_offline.py
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
  - pytest -v -s tests/distributed/test_context_parallel.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
+  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
+  - pytest -v -s tests/v1/distributed/test_dbo.py
+  - 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 (4xA100-4xMI300) # TBD
  timeout_in_minutes: 180
@@ -885,7 +715,7 @@ steps:
  commands:
  - pytest -v -s distributed/test_custom_all_reduce.py
  - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
-  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+  - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
  - pytest -v -s -x lora/test_mixtral.py

 - label: Distributed Torchrun + Examples (4 GPUs) # TBD
@@ -1369,7 +1199,7 @@ steps:
 #---------------------------------------------------------  mi300 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
@@ -1382,7 +1212,7 @@ steps:
  commands:
    - pip install tensorizer
    # Basic
-    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
+    - python3 basic/offline_inference/chat.py
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
    - python3 basic/offline_inference/classify.py
@@ -1586,10 +1416,10 @@ steps:
  - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset

 - label: Basic Models Tests (Other) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1631,6 +1461,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  torch_nightly: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
@@ -1835,6 +1666,54 @@ steps:
  - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
  - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper

+#----------------------------------------------------------  mi300 · plugins  ----------------------------------------------------------#
+
+- label: Plugin Tests (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/plugins/
+  - tests/plugins/
+  - vllm/platforms/rocm.py
+  commands:
+  # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform
+  - pip install -e ./plugins/vllm_add_dummy_platform
+  - pytest -v -s plugins_tests/test_platform_plugins.py
+  - pip uninstall vllm_add_dummy_platform -y
+  # END: platform plugin tests
+  # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin
+  - pip install -e ./plugins/prithvi_io_processor_plugin
+  - pytest -v -s plugins_tests/test_io_processor_plugins.py
+  - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
+  - pip uninstall prithvi_io_processor_plugin -y
+  # END: `io_processor` plugins test
+  # BEGIN: `bge_m3_sparse io_processor` test
+  - pip install -e ./plugins/bge_m3_sparse_plugin
+  - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
+  - pip uninstall bge_m3_sparse_plugin -y
+  # END: `bge_m3_sparse io_processor` test
+  # BEGIN: `colbert_query io_processor` test
+  - pip install -e ./plugins/colbert_query_plugin
+  - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
+  - pip uninstall colbert_query_plugin -y
+  # END: `colbert_query io_processor` test
+  # BEGIN: `stat_logger` plugins test
+  - pip install -e ./plugins/vllm_add_dummy_stat_logger
+  - pytest -v -s plugins_tests/test_stats_logger_plugins.py
+  - pip uninstall dummy_stat_logger -y
+  # END: `stat_logger` plugins test
+  # BEGIN: other tests
+  - pytest -v -s plugins_tests/test_scheduler_plugins.py
+  - pip install -e ./plugins/vllm_add_dummy_model
+  - pytest -v -s distributed/test_distributed_oot.py
+  - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
+  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
+  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#

 - label: Quantization # TBD
@@ -2197,6 +2076,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
+  optional: true
  num_gpus: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2215,45 +2095,11 @@ steps:
  - 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

- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - 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/spec_decode_acceptance_test.sh
-
- label: Distributed Tests (2xH100-2xMI300) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/v1/distributed/
-  - vllm/model_executor/layers/fused_moe/
-  - tests/v1/distributed/test_dbo.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
-  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
-  - pytest -v -s tests/v1/distributed/test_dbo.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
-  - pytest -v -s tests/distributed/test_packed_tensor.py
-
 - label: Metrics, Tracing (2 GPUs) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
+  optional: true
  num_gpus: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2279,11 +2125,28 @@ steps:
  commands:
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"

+- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/v1/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: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
  num_gpus: 4
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
@@ -2322,7 +2185,7 @@ steps:
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
  - tests/v1/kv_connector/nixl_integration/
  - vllm/platforms/rocm.py
  commands:
@@ -2334,9 +2197,10 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
  num_gpus: 4
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
  - tests/v1/kv_connector/nixl_integration/
  - vllm/platforms/rocm.py
  commands:
@@ -2348,6 +2212,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
  num_gpus: 4
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
@@ -2384,7 +2249,7 @@ steps:
 #------------------------------------------------------  mi300 · weight_loading  -------------------------------------------------------#

 - label: Weight Loading Multiple GPU # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -2396,7 +2261,7 @@ steps:
  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt

 - label: Weight Loading Multiple GPU - Large Models # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -2902,7 +2767,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/examples"
@@ -2914,7 +2779,7 @@ steps:
  commands:
  - pip install tensorizer
  # Basic
-  - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
+  - python3 basic/offline_inference/chat.py
  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
  - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
  - python3 basic/offline_inference/classify.py
@@ -3308,7 +3173,7 @@ steps:
 #------------------------------------------------------  mi355 · weight_loading  -------------------------------------------------------#

 - label: Weight Loading Multiple GPU # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_2
  num_gpus: 2
@@ -3320,7 +3185,7 @@ steps:
  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt

 - label: Weight Loading Multiple GPU - Large Models # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_2
  working_dir: "/vllm-workspace/tests"
@@ -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
@@ -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
@@ -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:
@@ -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
@@ -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/"
@@ -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
@@ -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
@@ -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
@@ -47,8 +47,10 @@ steps:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
  - vllm/models/deepseek_v4/common/ops/
  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
  commands:
    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py

 - label: Deepseek V4 Kernel Test (B200)
  key: deepseek-v4-kernel-test-b200
@@ -272,8 +274,8 @@ steps:
    - pytest -v -s kernels/helion/

 
- label: Kernels FP8 MoE Test (1 H100)
-  key: kernels-fp8-moe-test-1-h100
+- label: Kernels FP8 MoE Test (1xH100)
+  key: kernels-fp8-moe-test-1xh100
  timeout_in_minutes: 90
  device: h100
  num_devices: 1
@@ -289,8 +291,8 @@ steps:
    - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
    - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

- label: Kernels FP8 MoE Test (2 H100s)
-  key: kernels-fp8-moe-test-2-h100s
+- label: Kernels FP8 MoE Test (2xH100)
+  key: kernels-fp8-moe-test-2xh100
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
@@ -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
@@ -192,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
@@ -206,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
@@ -226,6 +239,8 @@ steps:
  device: dgx-spark
  optional: true
  num_devices: 1
+  depends_on:
+  - arm64-image-build
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -224,6 +224,16 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
+  mirror:
+    amd:
+      device: mi325_1
+      source_file_dependencies:
+      - vllm/entrypoints
+      - vllm/multimodal
+      - examples/
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd

 - label: Metrics, Tracing (2 GPUs)
  key: metrics-tracing-2-gpus
@@ -250,6 +260,12 @@ steps:
      'opentelemetry-exporter-otlp>=1.26.0' \
      'opentelemetry-semantic-conventions-ai>=0.4.1'"
  - pytest -v -s v1/tracing
+  mirror:
+    amd:
+      device: mi325_2
+      depends_on:
+      - image-build-amd
+      optional: true

 - label: Python-only Installation
  key: python-only-installation
@@ -262,6 +278,16 @@ steps:
  - setup.py
  commands:
  - bash standalone_tests/python_only_compile.sh
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 20
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/standalone_tests/python_only_compile.sh
+      - setup.py
+      - vllm/platforms/rocm.py

 - label: Async Engine, Inputs, Utils, Worker
  device: h200_35gb
@@ -344,7 +370,7 @@ steps:
  - pytest -v -s test_ray_env.py
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
+  - pytest -v -s reasoning
  - pytest -v -s tool_parsers
  - pytest -v -s tokenizers_
  - pytest -v -s parser
@@ -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
@@ -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
@@ -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
@@ -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
@@ -195,3 +195,11 @@ steps:
  - requirements/test/nightly-torch.txt
  commands:
  - bash standalone_tests/pytorch_nightly_dependency.sh
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - requirements/test/nightly-torch.txt
+      - vllm/platforms/rocm.py
@@ -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"
@@ -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"
@@ -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
@@ -3,3 +3,5 @@
 self-hosted-runner:
  labels:
    - vllm-runners
+    # Not yet in actionlint's known-label set.
+    - macos-26
@@ -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@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1

-      - 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"
@@ -270,6 +270,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 +392,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 +526,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 +879,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 +1085,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)
@@ -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)

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

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

@@ -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)
@@ -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
@@ -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;
@@ -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(
@@ -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);
@@ -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
@@ -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
@@ -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)));
@@ -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>(
@@ -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
@@ -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));
@@ -155,7 +155,7 @@ void cpu_gemm_wna16_impl(
  constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
  constexpr int32_t n_block_size = 16;
  static_assert(gemm_n_tile_size % n_block_size == 0);
-  const int32_t thread_num = omp_get_max_threads();
+  const int32_t thread_num = cpu_utils::get_max_threads();

  // a simple schedule policy, just to hold more B tiles in L2 and make sure
  // each thread has tasks
@@ -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) {
@@ -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) {
@@ -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) {
@@ -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
@@ -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) {
@@ -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));
@@ -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,
@@ -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]);
      }
    });
  }
@@ -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) {
@@ -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;
@@ -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,
@@ -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, "
@@ -2,13 +2,14 @@
 #define UTILS_HPP

 #include <atomic>
+#include <string>
 #include <unistd.h>
 #include <ATen/cpu/Utils.h>

 #include "cpu/cpu_types.hpp"

 namespace cpu_utils {
-enum class ISA { AMX, VEC, RVV };
+enum class ISA { AMX, VEC, RVV, NEON };

 inline ISA get_isa(const std::string& isa) {
  if (isa == "amx") {
@@ -17,6 +18,8 @@ inline ISA get_isa(const std::string& isa) {
    return ISA::VEC;
  } else if (isa == "rvv") {
    return ISA::RVV;
+  } else if (isa == "neon") {
+    return ISA::NEON;
  } else {
    TORCH_CHECK(false, "Invalid isa type: " + isa);
  }
@@ -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);

@@ -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;
-}
@@ -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
@@ -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
@@ -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
+}
@@ -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_
@@ -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());
@@ -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);
            }
          });
        });
@@ -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);
@@ -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);
    }
  }
 }
@@ -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);
@@ -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(
@@ -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,
@@ -11,6 +11,8 @@
 #include <cub/cub.cuh>
 #include <cstdint>

+#include "topk_histogram_4096.cuh"
+
 namespace vllm {
 namespace persistent {

@@ -935,8 +937,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
 }  // namespace persistent

 // ============================================================================
-// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only.
-// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent).
+// ============================================================================
+// Optimized FilteredTopK — single CTA per row for bs > 32.
+// Kept with persistent_topk so the portable fallback owns the non-cluster path.
+// ============================================================================
+namespace filtered_topk {
+
+namespace hist4096 = topk_histogram_4096;
+
+// ============================================================================
+// FilteredTopK — single CTA per row for bs > 32
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================

@@ -963,13 +973,6 @@ struct vec_t {
      data[i] = ptr[i];
    }
  }
-
-  FLASHINFER_INLINE void cast_store(T* ptr) const {
-#pragma unroll
-    for (size_t i = 0; i < N; ++i) {
-      ptr[i] = data[i];
-    }
-  }
 };
 #undef FLASHINFER_INLINE

@@ -1013,7 +1016,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
 * \tparam IdType Index type (int32_t)
 * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
 */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
+          bool UsePredicatedShortLoads = false>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                              IdType* __restrict__ output,
@@ -1042,6 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
    return;
  }

+  // Short path
+  if (length <= 32768) {
+    extern __shared__ uint8_t _smem_reg[];
+    if constexpr (UsePredicatedShortLoads) {
+      hist4096::histogram_4096_topk_predicated<MAX_K, 12, 8>(score, dst, length,
+                                                             _smem_reg);
+    } else {
+      hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
+                                                  _smem_reg);
+    }
+    return;
+  }
+
  // Static shared memory
  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
  alignas(128) __shared__ int s_counter;
@@ -1285,14 +1302,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,

  const int vec_size = ComputeFilteredTopKVecSize<DType>(max_len);

-#define DISPATCH_VEC_SIZE(VS)                                               \
-  if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
-    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
-        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
-    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
-                                          smem_size, stream));              \
-    return cudaSuccess;                                                     \
+#define DISPATCH_VEC_SIZE(VS)                                                 \
+  if (vec_size == VS) {                                                       \
+    auto kernel =                                                             \
+        FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K, (VS != MAX_VEC)>; \
+    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                                \
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));     \
+    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args,   \
+                                          smem_size, stream));                \
+    return cudaSuccess;                                                       \
  }

  DISPATCH_VEC_SIZE(1)
@@ -1306,6 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
  return cudaSuccess;
 }

+}  // namespace filtered_topk
+
+template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
+                                        IdType* output_indices,
+                                        const IdType* lengths,
+                                        uint32_t num_rows, uint32_t top_k_val,
+                                        uint32_t max_len,
+                                        cudaStream_t stream = 0) {
+  return filtered_topk::FilteredTopKRaggedTransform<DType, IdType, MAX_K>(
+      input, output_indices, lengths, num_rows, top_k_val, max_len, stream);
+}
+
 }  // namespace vllm

 #endif  // PERSISTENT_TOPK_CUH_
@@ -180,7 +180,7 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {

 #if defined CUDA_VERSION
  #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
-  if (cuda_device_capability >= 100 && cuda_device_capability < 110) {
+  if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
    return CUDA_VERSION >= 12080;
  }
  #endif
@@ -119,6 +119,10 @@ __global__ void per_token_group_quant_8bit_kernel(
      static_cast<DST_DTYPE*>(output_q) + block_group_offset;
  scale_element_t* scale_output;

+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaGridDependencySynchronize();
+#endif
+
  if constexpr (IS_COLUMN_MAJOR) {
    const int num_elems_per_pack =
        static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
@@ -153,6 +157,10 @@ __global__ void per_token_group_quant_8bit_kernel(

  QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
                              threads_per_group, y_s, min_8bit, max_8bit);
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }

 inline int GetGroupsPerBlock(int64_t num_groups) {
@@ -209,45 +217,56 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
  const int scale_num_rows = output_s.size(1);
  const int scale_stride = output_s.stride(1);

-#define LAUNCH_KERNEL(T, DST_DTYPE)                                        \
-  do {                                                                     \
-    dim3 grid(num_blocks);                                                 \
-    dim3 block(num_threads);                                               \
-    size_t smem_bytes =                                                    \
-        static_cast<size_t>(groups_per_block) * group_size * sizeof(T);    \
-    if (is_column_major) {                                                 \
-      if (scale_ue8m0) {                                                   \
-        per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, true>        \
-            <<<grid, block, smem_bytes, stream>>>(                         \
-                static_cast<T*>(input.data_ptr()), output_q.data_ptr(),    \
-                static_cast<float*>(output_s.data_ptr()), group_size,      \
-                num_groups, groups_per_block, (float)eps, (float)min_8bit, \
-                (float)max_8bit, scale_num_rows, scale_stride);            \
-      } else {                                                             \
-        per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, false>       \
-            <<<grid, block, smem_bytes, stream>>>(                         \
-                static_cast<T*>(input.data_ptr()), output_q.data_ptr(),    \
-                static_cast<float*>(output_s.data_ptr()), group_size,      \
-                num_groups, groups_per_block, (float)eps, (float)min_8bit, \
-                (float)max_8bit, scale_num_rows, scale_stride);            \
-      }                                                                    \
-    } else {                                                               \
-      if (scale_ue8m0) {                                                   \
-        per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, true>       \
-            <<<grid, block, smem_bytes, stream>>>(                         \
-                static_cast<T*>(input.data_ptr()), output_q.data_ptr(),    \
-                static_cast<float*>(output_s.data_ptr()), group_size,      \
-                num_groups, groups_per_block, (float)eps, (float)min_8bit, \
-                (float)max_8bit);                                          \
-      } else {                                                             \
-        per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, false>      \
-            <<<grid, block, smem_bytes, stream>>>(                         \
-                static_cast<T*>(input.data_ptr()), output_q.data_ptr(),    \
-                static_cast<float*>(output_s.data_ptr()), group_size,      \
-                num_groups, groups_per_block, (float)eps, (float)min_8bit, \
-                (float)max_8bit);                                          \
-      }                                                                    \
-    }                                                                      \
+#ifndef USE_ROCM
+  #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)     \
+    do {                                                                     \
+      cudaLaunchConfig_t config = {};                                        \
+      config.gridDim = dim3(num_blocks);                                     \
+      config.blockDim = dim3(num_threads);                                   \
+      config.dynamicSmemBytes = (SMEM_BYTES);                                \
+      config.stream = stream;                                                \
+      cudaLaunchAttribute attrs[1];                                          \
+      attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;      \
+      attrs[0].val.programmaticStreamSerializationAllowed = 1;               \
+      config.numAttrs = 1;                                                   \
+      config.attrs = attrs;                                                  \
+      cudaLaunchKernelEx(                                                    \
+          &config,                                                           \
+          per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
+          static_cast<T*>(input.data_ptr()), output_q.data_ptr(),            \
+          static_cast<float*>(output_s.data_ptr()), group_size, num_groups,  \
+          groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit,    \
+          scale_num_rows, scale_stride);                                     \
+    } while (0)
+#else
+  #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)   \
+    do {                                                                   \
+      per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>    \
+          <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
+              static_cast<T*>(input.data_ptr()), output_q.data_ptr(),      \
+              static_cast<float*>(output_s.data_ptr()), group_size,        \
+              num_groups, groups_per_block, (float)eps, (float)min_8bit,   \
+              (float)max_8bit, scale_num_rows, scale_stride);              \
+    } while (0)
+#endif
+
+#define LAUNCH_KERNEL(T, DST_DTYPE)                                     \
+  do {                                                                  \
+    size_t smem_bytes =                                                 \
+        static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
+    if (is_column_major) {                                              \
+      if (scale_ue8m0) {                                                \
+        LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes);       \
+      } else {                                                          \
+        LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes);      \
+      }                                                                 \
+    } else {                                                            \
+      if (scale_ue8m0) {                                                \
+        LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes);      \
+      } else {                                                          \
+        LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes);     \
+      }                                                                 \
+    }                                                                   \
  } while (0)

  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
@@ -262,6 +281,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
      }));

 #undef LAUNCH_KERNEL
+#undef LAUNCH_KERNEL_INST
 }

 // Register-resident fast path for group_size==128.
@@ -306,12 +326,12 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
  const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;

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

  if (mn_idx >= tma_aligned_mn) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
    return;
  }
@@ -428,7 +448,7 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
  *reinterpret_cast<uint4*>(group_output) = packed_out;

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

@@ -0,0 +1,563 @@
+/*
+ * Shared 4096-bin single-CTA TopK helpers.
+ */
+
+#ifndef TOPK_HISTOGRAM_4096_CUH_
+#define TOPK_HISTOGRAM_4096_CUH_
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace vllm {
+namespace topk_histogram_4096 {
+
+constexpr uint32_t kBlockSize = 1024;
+constexpr uint32_t RADIX = 256;
+constexpr uint32_t kMaxTies = 1024;
+static_assert(kMaxTies <= kBlockSize,
+              "tie_handle requires kMaxTies <= kBlockSize");
+constexpr uint32_t kWarpSize = 32;
+constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
+
+// Register path
+constexpr uint32_t kHist4096VecsPerThread = 4;
+constexpr uint32_t kHist4096MaxLen =
+    kHist4096VecsPerThread * 4 * kBlockSize;  // 16384
+
+struct alignas(16) MatchBin {
+  uint32_t bin, above_count, equal_count;
+};
+struct alignas(8) Tie {
+  uint32_t idx;
+  float score;
+};
+
+__device__ __forceinline__ void load_float4_predicated(const float* ptr,
+                                                       int base, int seq_len,
+                                                       float& v0, float& v1,
+                                                       float& v2, float& v3) {
+  uint32_t r0, r1, r2, r3;
+  const int p0 = (base < seq_len);
+  const int p1 = (base + 1 < seq_len);
+  const int p2 = (base + 2 < seq_len);
+  const int p3 = (base + 3 < seq_len);
+  asm volatile(
+      "{\n"
+      "  .reg .pred pr0, pr1, pr2, pr3;\n"
+      "  setp.ne.u32 pr0, %4, 0;\n"
+      "  setp.ne.u32 pr1, %5, 0;\n"
+      "  setp.ne.u32 pr2, %6, 0;\n"
+      "  setp.ne.u32 pr3, %7, 0;\n"
+      "  mov.u32 %0, 0xFF800000;\n"
+      "  mov.u32 %1, 0xFF800000;\n"
+      "  mov.u32 %2, 0xFF800000;\n"
+      "  mov.u32 %3, 0xFF800000;\n"
+      "  @pr0 ld.global.cg.u32 %0, [%8];\n"
+      "  @pr1 ld.global.cg.u32 %1, [%8+4];\n"
+      "  @pr2 ld.global.cg.u32 %2, [%8+8];\n"
+      "  @pr3 ld.global.cg.u32 %3, [%8+12];\n"
+      "}\n"
+      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+      : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr));
+  v0 = __uint_as_float(r0);
+  v1 = __uint_as_float(r1);
+  v2 = __uint_as_float(r2);
+  v3 = __uint_as_float(r3);
+}
+
+// converts the float32 score to a 32-bit ordered unsigned integer — the full
+// precision key for radix sorting
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits ->
+// bin 0-4095)
+template <uint32_t kBits>
+__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(bits | 0x8000);
+  return key >> (16 - kBits);
+}
+
+// running sum within each warp — thread 0 gets its own value, thread 1 gets
+// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc.
+__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane,
+                                                       uint32_t v) {
+#pragma unroll
+  for (uint32_t o = 1; o < 32; o *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o);
+    if (lane >= o) v += n;
+  }
+  return v;
+}
+
+// Returns the sum of a value across all 32 threads in the warp, and every
+// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX
+// instruction for hardware warp-wide reduction. Older targets use the
+// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32
+// lanes).
+__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  uint32_t r;
+  asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v));
+  return r;
+#else
+  #pragma unroll
+  for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) {
+    v += __shfl_xor_sync(0xFFFFFFFF, v, mask);
+  }
+  return v;
+#endif
+}
+
+// ============================================================================
+// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered
+// key Each round narrows by 8 bits until ties are fully resolved
+// ============================================================================
+
+template <uint32_t TopK>
+__device__ void tie_handle(const Tie* ties, uint32_t num_ties,
+                           uint32_t num_above, int32_t* output, void* _smem) {
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize, wi = tx / kWarpSize;
+
+  // Each thread loads one tie element.
+  const bool has = tx < num_ties;
+  const auto tie = has ? ties[tx] : Tie{0, 0.0f};
+  const uint32_t key = convert_to_uint32_v2(tie.score);
+
+  bool active = has;  // tracks whether this thread's tie is still a candidate.
+  uint32_t remain =
+      TopK - num_above;  // decreases each round as ties are resolved.
+  uint32_t wpos = TopK;  // wpos will hold the final output position.
+  s->counter = 0;
+  __syncthreads();
+
+  // The 4-round radix loop - each round narrows by 8 bits until ties are fully
+  // resolved
+#pragma unroll
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;  // round 0: bits 31-24, round 1: 23-16, etc.
+    uint32_t bin = (key >> sh) & 0xFF;  // this tie's 8-bit bin for this round
+
+    // Step 1: Build 256-bin histogram.
+    if (tx < RADIX) s->histogram[tx] = 0;
+    __syncthreads();
+    if (active) atomicAdd(&s->histogram[bin], 1);
+    __syncthreads();
+
+    // Step 2: Prefix scan to find threshold
+    uint32_t hv = 0, wi2 = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) s->warp_sum[wi] = wi2;
+    }
+    __syncthreads();
+
+    if (tx < RADIX) {
+      auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto tot = warp_reduce_sum_full(tmp);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp : 0);
+      auto above = tot - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {tx, above, remain - above};
+      }
+    }
+    __syncthreads();
+
+    // Step 3: Scatter
+    auto [thr, na, _] = s->match;  // threshold bin, num above, unused
+    if (active) {
+      if (bin > thr) {
+        wpos = num_above +
+               atomicAdd(&s->counter, 1);  // above -> place in output directly
+        active = false;
+      } else if (bin < thr)
+        active = false;  // below -> discard
+      else if (r == 3)
+        wpos = TopK - atomicAdd(&s->match.equal_count,
+                                -1u);  // last round: place remaining
+    }
+    remain -= na;
+    if (!remain) break;  // all ties resolved early
+  }
+  // Final write
+  if (wpos < TopK) output[wpos] = tie.idx;
+}
+
+// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048).
+// tie_handle assumes 1 tie per thread (max 1024).
+// This version handles 2 ties per thread via kPerThread=2
+template <uint32_t TopK>
+__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties,
+                                 uint32_t num_above, int32_t* output,
+                                 void* _smem) {
+  static_assert(TopK > kBlockSize);
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize;
+  const auto wi = tx / kWarpSize;
+
+  constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize;
+  Tie my_ties[kPerThread];
+  uint32_t keys[kPerThread];
+  bool active[kPerThread];
+
+  for (uint32_t e = 0; e < kPerThread; e++) {
+    uint32_t idx = e * kBlockSize + tx;
+    if (idx < num_ties) {
+      my_ties[e] = ties[idx];
+      keys[e] = convert_to_uint32_v2(ties[idx].score);
+      active[e] = true;
+    } else {
+      my_ties[e] = {0, 0.0f};
+      keys[e] = 0;
+      active[e] = false;
+    }
+  }
+
+  uint32_t remain = TopK - num_above;
+  s->counter = 0;
+  __syncthreads();
+
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;
+    if (tx < RADIX) {
+      s->histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (active[e]) {
+        atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1);
+      }
+    }
+    __syncthreads();
+
+    uint32_t hv = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      auto wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) {
+        s->warp_sum[wi] = wi2;
+      }
+    }
+    __syncthreads();
+    if (tx < RADIX) {
+      auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto total = warp_reduce_sum_full(tmp2);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0);
+      auto wi2 = warp_inclusive_sum(li, hv);
+      auto above = total - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {
+            .bin = tx, .above_count = above, .equal_count = remain - above};
+      }
+    }
+    __syncthreads();
+
+    auto thr = s->match.bin;
+    auto na = s->match.above_count;
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (!active[e]) {
+        continue;
+      }
+      uint32_t bin = (keys[e] >> sh) & 0xFF;
+      if (bin > thr) {
+        uint32_t wpos = num_above + atomicAdd(&s->counter, 1);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+        active[e] = false;
+      } else if (bin < thr) {
+        active[e] = false;
+      } else if (r == 3) {
+        uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+      }
+    }
+
+    num_above += na;
+    remain -= na;
+    __syncthreads();
+    s->counter = 0;
+    __syncthreads();
+  }
+}
+
+// ============================================================================
+// Register-based single-CTA fast path for seq_len <= 16384
+// 4 float4 per thread × 1024 threads = 16384 elements max
+// Uses 4096-bin (12-bit) histogram for better precision
+// ============================================================================
+
+template <uint32_t TopK, uint32_t HIST_BITS>
+struct Histogram4096Smem {
+  static constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies;
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  MatchBin match;
+  uint32_t warp_sum[kNumWarps];
+  union {
+    uint32_t histogram[HIST_BINS];
+    Tie tie_buffer[TIE_CAPACITY];
+  };
+};
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
+          bool UsePredicatedLoads = false>
+__device__ void histogram_4096_topk(const float* __restrict__ scores,
+                                    int32_t* __restrict__ output,
+                                    uint32_t length, void* _smem) {
+  constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize;
+  static_assert(HIST_BINS >= kBlockSize,
+                "HIST_BITS must give >= kBlockSize bins");
+
+  using Smem = Histogram4096Smem<TopK, HIST_BITS>;
+  auto* smem = static_cast<Smem*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto lane_id = tx % kWarpSize;
+  const auto warp_id = tx / kWarpSize;
+
+  // Phase 1: Load all data into RF + build histogram
+  float4
+      vecs[VECS_PER_THREAD];  // 4 vectors x 4 floats = 16 elements per thread
+  if constexpr (ITEMS_PER_THREAD >= 4) {
+    // Zero the histogram (SMEM writes)
+    for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++)
+      reinterpret_cast<uint4*>(
+          smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] =
+          make_uint4(0, 0, 0, 0);
+  } else {
+    if (tx < HIST_BINS) smem->histogram[tx] = 0;
+  }
+  if (tx == 0) {
+    smem->counter_gt = 0;
+    smem->counter_eq = 0;
+  }
+  if constexpr (UsePredicatedLoads) {
+    const bool row_aligned = (reinterpret_cast<uintptr_t>(scores) & 0xFu) == 0;
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        if (row_aligned && base + 3 < length) {
+          vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+        } else {
+          load_float4_predicated(scores + base, static_cast<int>(base),
+                                 static_cast<int>(length), vecs[v].x, vecs[v].y,
+                                 vecs[v].z, vecs[v].w);
+        }
+      }
+    }
+  } else {
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Build histogram from RF via atomic adds into the shared histogram
+  bool done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        atomicAdd(&smem->histogram[extract_coarse_bin_N<HIST_BITS>(elems[e])],
+                  1);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Phase 2: Prefix scan to find threshold bin
+  // Multi-element scan (4096 bins: 4 per thread)
+  uint32_t orig[ITEMS_PER_THREAD];
+  uint32_t local_sum = 0;
+
+  // Step 1: Each thread sums its 4 bins
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i];
+    local_sum += orig[i];
+  }
+
+  // Step 2: Warp-level inclusive prefix sum on local_sum
+  const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
+  if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+  __syncthreads();
+
+  // Step 3: Inter-warp prefix across warp sums.
+  const auto tmp = smem->warp_sum[lane_id];
+  uint32_t prefix = warp_reduce_sum_full(
+      lane_id < warp_id ? tmp : 0);  // sum of all prior warps
+  prefix +=
+      warp_inc - local_sum;  // exclusive prefix within this thread's position
+
+  // Step 4: Find threshold - scan 4 bins, accumulate prefix
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    prefix += orig[i];
+    const auto above = length - prefix;  // elements in bins ABOVE this one
+    if (above < TopK && above + orig[i] >= TopK) {
+      smem->match = {.bin = tx * ITEMS_PER_THREAD + i,
+                     .above_count = above,
+                     .equal_count = orig[i]};
+    }
+  }
+
+  __syncthreads();
+
+  // Phase 3: Scatter from registers
+  const auto [thr_bin, num_above, num_equal] = smem->match;
+  const bool need_tie = (num_equal + num_above > TopK);
+
+  done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        const uint32_t bin = extract_coarse_bin_N<HIST_BITS>(elems[e]);
+        if (bin > thr_bin) {
+          output[atomicAdd(&smem->counter_gt, 1)] =
+              idx;  // above -> output directly
+        } else if (bin == thr_bin) {
+          const auto pos = atomicAdd(&smem->counter_eq, 1);
+          if (!need_tie) {
+            if (pos + num_above < TopK) {
+              output[pos + num_above] = idx;  // all fit
+            }
+          } else {
+            if (pos < TopK) {
+              smem->tie_buffer[pos] = {idx, elems[e]};  // store for refirement
+            }
+          }
+        }
+        // else: bin < thr_bin - discard (not in top-k)
+      }
+    }
+  }
+
+  // Phase 4: Tie-breaking
+  if (!need_tie) return;
+  __syncthreads();
+
+  // Fast warp-ballot tie-breaking for small tie counts
+  const uint32_t num_ties = min(num_equal, static_cast<uint32_t>(TopK));
+  const uint32_t topk_remain =
+      TopK - num_above;  // pick exactly remaining elements to fill topK
+
+  auto is_greater = [](const Tie& a, const Tie& b) {
+    return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
+  };
+
+  if (num_ties <= kWarpSize) {
+    // <=32 ties - Use warp ballot
+    // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024
+    // comparisons in one instruction per warp. O(1) work.
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    if (lane_id >= num_ties || warp_id >= num_ties) return;
+    const uint32_t mask = (1ull << num_ties) - 1u;
+    const auto tie = smem->tie_buffer[lane_id];  // each lane holds one tie
+    const auto target =
+        smem->tie_buffer[warp_id];  // each warp evaluates one candidate
+    const bool pred =
+        is_greater(tie, target);  // compare all ties against target
+    const auto rank = static_cast<uint32_t>(
+        __popc(__ballot_sync(mask, pred)));  // count how many are greater
+    if (lane_id == 0 && rank < topk_remain) {
+      output[num_above + rank] = target.idx;  // place at correct position
+    }
+  } else if (num_ties <=
+             kWarpSize *
+                 2) {  // TODO (roberto): try to refactor this with <=32 case
+    //  Same idea but each thread handles 2 tie elements
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    const auto lane1 = lane_id + kWarpSize;
+    const auto warp1 = warp_id + kWarpSize;
+    const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__};
+    const auto tie0 = smem->tie_buffer[lane_id];
+    const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid;
+    if (warp_id < num_ties) {
+      const auto target = smem->tie_buffer[warp_id];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+    if (warp1 < num_ties) {
+      const auto target = smem->tie_buffer[warp1];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+  } else {
+    // Large tie count: fall back to 4-round radix-256 sort
+    if constexpr (TopK <= kBlockSize) {
+      tie_handle<TopK>(smem->tie_buffer, num_ties, num_above, output, smem);
+    } else {
+      tie_handle_large<TopK>(smem->tie_buffer, num_ties, num_above, output,
+                             smem);
+    }
+  }
+}
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
+__device__ __noinline__ void histogram_4096_topk_predicated(
+    const float* __restrict__ scores, int32_t* __restrict__ output,
+    uint32_t length, void* _smem) {
+  histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
+                                                              length, _smem);
+}
+
+}  // namespace topk_histogram_4096
+}  // namespace vllm
+
+#endif  // TOPK_HISTOGRAM_4096_CUH_
@@ -493,6 +493,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
      "Tensor workspace, int k, int max_seq_len) -> ()");

+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.def(
+      "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, "
+      "Tensor workspace, int k, int max_seq_len) -> ()");
+#endif
+
  // Activation ops
  ops.def(
      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
@@ -711,6 +717,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
  ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
  ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk));
+#endif

  // Activation kernels (shared CUDA/ROCm)
  ops.impl("persistent_masked_m_silu_mul_quant",
@@ -41,12 +41,6 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);

 void gelu_quick(torch::Tensor& out, torch::Tensor& input);

-void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
-                        torch::Tensor const& q_pe,
-                        torch::Tensor const& kv_c_and_k_pe_cache,
-                        torch::Tensor const& seq_lens,
-                        torch::Tensor const& page_table, double scale);
-
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor const& scale,
                              std::optional<torch::Tensor> const& azp);
@@ -283,6 +283,29 @@ __quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
  return R.i;
 }

+template <typename T>
+__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low,
+                                                         int16_t high);
+
+template <>
+__quickreduce_device_inline__ int packed_from_int16_pair<half>(int16_t low,
+                                                               int16_t high) {
+  // Convert two signed integers to one fp16x2 packed 32-bit lane.
+  half2 h = __halves2half2(__int2half_rn(static_cast<int>(low)),
+                           __int2half_rn(static_cast<int>(high)));
+  return __builtin_bit_cast(int, h);
+}
+
+template <>
+__quickreduce_device_inline__ int packed_from_int16_pair<nv_bfloat16>(
+    int16_t low, int16_t high) {
+  // Convert two signed integers to one bf16x2 packed 32-bit lane.
+  nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
+  nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
+  nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
+  return *reinterpret_cast<int*>(&bf2);
+}
+
 // changes dtype
 __quickreduce_device_inline__ float T2float_cast(half a) {
  return __half2float(a);
@@ -59,11 +59,30 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
                       flag_color, this->kMaxProblemSize);                  \
  }

+// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
+// the 3-bit codec's pack/unpack overhead outweighs the reduced communication
+// volume, so INT3 is restricted to a TP2-only dispatch here.
+#define TWOSHOT_DISPATCH_TP2_ONLY(__codec)                                  \
+  if (world_size == 2) {                                                    \
+    using LineCodec = __codec<T, 2>;                                        \
+    using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
+    hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
+                       dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
+                       num_blocks, rank, dbuffer_list, data_offset,         \
+                       flag_color, this->kMaxProblemSize);                  \
+  } else {                                                                  \
+    throw std::runtime_error(                                               \
+        "INT3 quick all-reduce is only supported for world_size == 2 "      \
+        "(TP2); use INT4/NONE for larger world sizes.");                    \
+  }
+
 enum QuickReduceQuantLevel {
-  F16 = 0,
-  INT8 = 1,
-  INT6 = 2,
-  INT4 = 3,
+  // Keep these ids in sync with Python QuickReduceRegime enum.
+  F16 = 0,   // full-precision fp16/bf16 communication
+  INT8 = 1,  // symmetric int8 + per-block scale
+  INT6 = 2,  // symmetric int6 + per-block scale
+  INT4 = 3,  // symmetric int4 + per-block scale
+  INT3 = 4,  // symmetric int3 + per-block scale (TP2 only)
 };

 struct DeviceComms {
@@ -184,6 +203,9 @@ struct DeviceComms {
      case QuickReduceQuantLevel::INT4:
        TWOSHOT_DISPATCH(CodecQ4)
        break;
+      case QuickReduceQuantLevel::INT3:
+        TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3)
+        break;
      default:
        TWOSHOT_DISPATCH(CodecFP)
        break;
@@ -2,6 +2,7 @@

 #include <hip/hip_runtime.h>
 #include "base.h"
+#include <hip/hip_fp16.h>

 namespace quickreduce {

@@ -206,6 +207,168 @@ struct CodecQ4 : public CodecBase {
  }
 };

+// Int3 symmetric quantization codec.
+// We quantize the FP16 data to block-scaled Int3 in blocks of 4 *
+// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip),
+// matching the structure of CodecQ4. Signed range is [-4, +3].
+template <typename T, int world_size>
+struct CodecQ3 : public CodecBase {
+  static constexpr int kWorldSize = world_size;
+
+  // Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes):
+  // - each thread owns 8 values and writes:
+  //   * q2 payload  : 8 * 2 bits  -> uint16 (2 bytes)
+  //   * q1 payload  : 8 * 1 bit   -> uint8  (1 byte)
+  // - one scale is shared per 32 values and written by group leader.
+  //
+  // kRankTileStride is split as:
+  // [0 .. 511]   : q2 payload region   (256 threads * 2 bytes)
+  // [512 .. 767] : q1 payload region   (256 threads * 1 byte)
+  // [768 .. 895] : scale region        (32 groups   * 4 bytes)
+  static constexpr int kRankAtoms = kAtoms / kWorldSize;
+  static constexpr int kRankTileStride = 896;
+  static constexpr int kRankTileQ1Offset = 512;
+  static constexpr int kRankTileScaleOffset = 768;
+  static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
+  static_assert(kRankTransmittedTileSize % 16 == 0,
+                "kRankTransmittedTileSize must be 16B aligned.");
+
+  static constexpr int kRankBufferTileStride =
+      kRankTileStride / sizeof(int32x4_t);
+
+  static constexpr int kTransmittedTileSize =
+      kRankTransmittedTileSize * kWorldSize;
+
+  // {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps
+  // to -4; the sign cancels with decoding_scale on the recv side.
+  static constexpr int kScaleFactor =
+      std::is_same<T, half>::value ? 0xB400B400 : 0xBE80BE80;
+
+  // {1e-7, 1e-7}, f16x2_t
+  static constexpr int kScaleEpsilon =
+      std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
+
+  // {-4, -4}, f16x2_t / bf16x2_t
+  static constexpr int kRangeMin =
+      std::is_same<T, half>::value ? 0xC400C400 : 0xC080C080;
+
+  // {+3, +3}, f16x2_t / bf16x2_t
+  static constexpr int kRangeMax =
+      std::is_same<T, half>::value ? 0x42004200 : 0x40404040;
+
+  // {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7].
+  static constexpr int kRangeBias = 0x00040004;
+
+  __quickreduce_device_inline__ CodecQ3(int thread, int rank)
+      : CodecBase(thread, rank) {}
+
+  __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
+                                          const int32x4_t* __restrict__ data) {
+    for (int k = 0; k < kRankAtoms; k++) {
+      int32x4_t const atom = data[k];
+
+      // 1) Per-group dynamic scale (shared across 32 values).
+      int wblockmax = group_abs_max<T>(atom);
+      int decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
+      int encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
+      encoding_scale = packed_rcp<T>(encoding_scale);
+
+      // 2) Scale + clip to signed int3 range [-4, +3].
+      int32x4_t w;
+      for (int i = 0; i < 4; i++) {
+        w[i] = packed_mul<T>(atom[i], encoding_scale);
+        w[i] = packed_max<T>(w[i], kRangeMin);
+        w[i] = packed_min<T>(w[i], kRangeMax);
+      }
+
+      // 3) Round to integer and bias to unsigned domain [0, 7].
+      int32x4_t q;
+      {
+        int16_t* qi = reinterpret_cast<int16_t*>(&q);
+        T* wh = reinterpret_cast<T*>(&w);
+        for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
+
+        for (int i = 0; i < 4; i++) {
+          q[i] = packed_add<int16_t>(q[i], kRangeBias);
+        }
+      }
+
+      // 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit
+      //    halves, packed into one uint16 (low 2 bits per value) plus one
+      //    uint8 (high 1 bit per value).
+      uint16_t q2w = 0;
+      uint8_t q1w = 0;
+      {
+        int16_t* tw = reinterpret_cast<int16_t*>(&q);
+#pragma unroll
+        for (int i = 0; i < 8; i++) {
+          uint32_t v = static_cast<uint32_t>(tw[i]) & 0x7u;
+          q2w |= static_cast<uint16_t>((v & 0x3u) << (i * 2));
+          q1w |= static_cast<uint8_t>(((v >> 2) & 0x1u) << i);
+        }
+      }
+
+      uint8_t* atom_ptr =
+          reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
+      uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
+      uint8_t* q1w_ptr =
+          reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
+      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
+                    (thread / 8);
+
+      __builtin_nontemporal_store(q2w, q2w_ptr);
+      *q1w_ptr = q1w;
+      if (threadIdx.x == group_leader) {
+        __builtin_nontemporal_store(decoding_scale, qs_ptr);
+      }
+    }
+  }
+
+  __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
+                                          int32x4_t* __restrict__ data) {
+    for (int k = 0; k < kRankAtoms; k++) {
+      uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
+      uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
+      uint8_t* q1w_ptr =
+          reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
+      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
+                    (thread / 8);
+
+      uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
+      uint8_t q1w = *q1w_ptr;
+      int qs = __builtin_nontemporal_load(qs_ptr);
+
+      *recv_buffer += kRankBufferTileStride;
+
+      // Unpack unsigned values [0, 7] then shift back to signed domain
+      // [-4, +3] by adding kRangeMin.
+      int32x4_t w;
+      {
+        int16_t qv[8];
+#pragma unroll
+        for (int i = 0; i < 8; i++) {
+          uint32_t low2 = (q2w >> (2 * i)) & 0x3u;
+          uint32_t high1 = (q1w >> i) & 0x1u;
+          qv[i] = static_cast<int16_t>(low2 | (high1 << 2));
+        }
+
+#pragma unroll
+        for (int i = 0; i < 4; i++) {
+          int qpack = packed_from_int16_pair<T>(qv[2 * i], qv[2 * i + 1]);
+          w[i] = packed_add<T>(qpack, kRangeMin);
+        }
+      }
+
+      // Apply decode scale to reconstruct fp16/bf16 lanes.
+      for (int i = 0; i < 4; i++) {
+        w[i] = packed_mul<T>(w[i], qs);
+      }
+
+      data[k] = w;
+    }
+  }
+};
+
 // Int6 symmetric quantization codec.
 // We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
 // kThreadGroupSize.
@@ -377,7 +540,6 @@ struct CodecQ6 : public CodecBase {
        w[i] = packed_mul<T>(w[i], qs);
      }

-      // That's pretty much it...
      data[k] = w;
    }
  }
@@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
    const scalar_t* q_ptr =
        q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE;
    const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);
-    const int qhead_elemh8 = laneid / 4;
+    const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1);

    for (int h = 0; h < QHLOOP - 1; h++) {
      const int qhead_idx = h * 4 + lane4id;
@@ -70,6 +70,15 @@ bool on_gfx12() {
  return result;
 }

+bool on_gfx1151() {
+  static const bool result = [] {
+    const auto* dprops = at::cuda::getCurrentDeviceProperties();
+    const std::string device_arch = dprops->gcnArchName;
+    return device_arch.find("gfx1151") != std::string::npos;
+  }();
+  return result;
+}
+
 #if defined(NDEBUG)
  #undef NDEBUG
  #include <assert.h>
@@ -1237,6 +1246,45 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
      WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N)           \
  }

+// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N).
+//   THRDS  = wavefront width (32 on GFX11/GFX12, 64 on GFX9)
+//   WVPRGRP= waves per group (always 16)
+//   YTILE  = output rows per thread tile
+//   UNRL   = K-loop unroll factor
+//   N      = batch size (passed through from the switch in wvSplitK)
+#define WVSPLIT_TILE(_sYT, __N)                                             \
+  {                                                                         \
+    if (on_gfx1151()) {                                                     \
+      bool fit_lds = (Kbp_in * N_in <= max_lds_len);                        \
+      if (_sYT <= 1)                                                        \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
+                     __N)                                                   \
+      else if ((K_in % 1024 == 512) && K_in >= 1536 &&                      \
+               (_sYT >= 40 || K_in >= 4096))                                \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/4, /*UNRL=*/1, \
+                     __N)                                                   \
+      else if (K_in < 1024)                                                 \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/2, /*UNRL=*/4, \
+                     __N)                                                   \
+      else if (K_in <= 2048 && (__N >= 2 || _sYT <= 26))                    \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
+                     __N)                                                   \
+      else if (__N >= 2 && !fit_lds)                                        \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
+                     __N)                                                   \
+      else if (__N == 1)                                                    \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/2, \
+                     __N)                                                   \
+      else                                                                  \
+        WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/1, \
+                     __N)                                                   \
+    } else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */            \
+      WVSPLIT_TILE_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, _sYT, __N)             \
+    } else { /* GFX9, wave64 */                                             \
+      WVSPLIT_TILE_CFG(/*THRDS=*/64, /*WVPRGRP=*/16, _sYT, __N)             \
+    }                                                                       \
+  }
+
  AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
    using fptype = typename scalar<scalar_t>::type;
    fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
@@ -1251,37 +1299,21 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
    // then cut the active waves to balance their distribution...
    int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);

-    const bool use_wave32 = on_gfx1x();
    switch (N_in) {
      case 1:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 1)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 1)
+        WVSPLIT_TILE(sYT, 1)
        break;
      case 2:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 2)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 2)
+        WVSPLIT_TILE(sYT, 2)
        break;
      case 3:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 3)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 3)
+        WVSPLIT_TILE(sYT, 3)
        break;
      case 4:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 4)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 4)
+        WVSPLIT_TILE(sYT, 4)
        break;
      case 5:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 5)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 5)
+        WVSPLIT_TILE(sYT, 5)
        break;
      default:
        throw std::runtime_error(
@@ -311,9 +311,10 @@ ENV CARGO_BUILD_JOBS=4
 # Build the release artifacts. Cache cargo registry/git, but not target/,
 # because stale target metadata can outlive source updates across BuildKit
 # cache reuse.
-RUN --mount=type=cache,target=/root/.cargo/registry \
-    --mount=type=cache,target=/root/.cargo/git \
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
    bash build_rust.sh
+
 #################### RUST BUILD IMAGE ####################

 #################### CSRC BUILD IMAGE ####################
@@ -898,13 +899,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
        fi; \
    fi

-# Download FlashInfer precompiled cubins AFTER all pip installs are done.
-# This must run after the vLLM wheel and EP kernels installs above, because
-# those can reinstall/touch flashinfer packages. Downloading cubins earlier
-# (in the flashinfer-jit-cache layer) causes ~2.5 GB of layer duplication
-# when a later pip install overwrites flashinfer package files.
-RUN flashinfer show-config && flashinfer download-cubin
-
 # CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
 # return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
 # consistently from the host (see https://github.com/vllm-project/vllm/issues/18859).
@@ -123,9 +123,10 @@ COPY build_rust.sh build_rust.sh
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4

-RUN --mount=type=cache,target=/root/.cargo/registry \
-    --mount=type=cache,target=/root/.cargo/git \
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
    bash build_rust.sh
+
 #################### RUST BUILD IMAGE ####################

 #################### WHEEL BUILD IMAGE ####################
@@ -575,6 +575,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV HF_XET_HIGH_PERFORMANCE=1
 ENV HF_HUB_DOWNLOAD_TIMEOUT=60

+# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
+ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
+
 # Pre-install vLLM test dependencies.
 COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -695,6 +698,9 @@ ENV SAFETENSORS_FAST_GPU=1
 # Performance environment variable.
 ENV HIP_FORCE_DEV_KERNARG=1

+# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
+ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
+
 # Workaround for ROCm profiler limits
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
 ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.13.post1"
+ARG AITER_BRANCH="v0.1.16.post2"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -244,7 +244,7 @@ RUN pip install pyyaml && cd aiter \
           export HIP_CLANG_PATH=/opt/sccache-wrappers \
           && sccache --show-stats; \
       fi \
-    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
    && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
@@ -28,22 +28,26 @@ COPY build_rust.sh build_rust.sh
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4

-RUN --mount=type=cache,target=/root/.cargo/registry \
-    --mount=type=cache,target=/root/.cargo/git \
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
    bash build_rust.sh

-FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
+FROM ubuntu:24.04 AS vllm-base
+
+ENV DEBIAN_FRONTEND=noninteractive

 WORKDIR /workspace/

 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"

-RUN apt clean && apt-get update -y && \
-    apt-get install -y --no-install-recommends --fix-missing \
+RUN apt-get update -y && \
+    apt-get install -y --no-install-recommends \
+    build-essential \
    curl \
    ffmpeg \
    git \
+    gpg \
    libsndfile1 \
    libsm6 \
    libxext6 \
@@ -53,9 +57,11 @@ RUN apt clean && apt-get update -y && \
    numactl \
    wget \
    vim \
+    ca-certificates \
    python3.12 \
    python3.12-dev \
-    python3-pip
+    python3-pip && \
+    rm -rf /var/lib/apt/lists/*

 # Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
 RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
@@ -82,6 +88,7 @@ RUN mkdir neo && \
    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
    wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero-devel_1.28.2+u24.04_amd64.deb && \
    dpkg -i *.deb && \
    cd .. && \
    rm -rf neo
@@ -101,7 +108,13 @@ RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${
    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
    rm -f /opt/intel/oneapi/ccl/latest && \
-    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \
+    printf '%s\n' \
+    '/opt/intel/oneapi/ccl/2021.15/lib' \
+    '/opt/intel/oneapi/mpi/2021.15/lib' \
+    '/opt/intel/oneapi/compiler/2025.3/lib' \
+    > /etc/ld.so.conf.d/oneapi-ccl.conf && \
+    ldconfig

 SHELL ["bash", "-c"]
 CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
@@ -123,7 +136,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \



-ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib
 CMD ["/bin/bash"]

 ######################### UCX + NIXL BUILD STAGE #########################
@@ -204,10 +217,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
    uv pip uninstall triton triton-xpu && \
    uv pip install triton-xpu==3.7.1 && \
-    uv pip uninstall oneccl oneccl-devel && \
-    source /opt/intel/oneapi/setvars.sh --force && \
-    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
-    export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
+    uv pip uninstall oneccl oneccl-devel

 # Keep source-dependent layers near the end so frequent code-only changes
 # don't invalidate heavy dependency and UCX/NIXL layers.
@@ -302,14 +302,22 @@ group "test-rocm-ci-with-wheel" {
 }

 # Image tags for the ci_base build. ci-bake-rocm.sh rewrites CI_BASE_IMAGE_TAG
-# to the primary tag for this build. Non-nightly builds use a commit-scoped tag
-# and also publish a content tag for reuse. NIGHTLY=1 builds on the stable branch
-# can additionally set CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base.
+# to the primary tag for this build. Builds always publish a content-scoped tag
+# when the ci_base content hash is available. Builds with BUILDKITE_COMMIT also
+# publish a commit-scoped tag, either as the primary tag or an additional alias.
+# NIGHTLY=1 builds on the stable branch can additionally set
+# CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base.
 variable "CI_BASE_IMAGE_TAG" {
  default = "rocm/vllm-dev:ci_base"
 }

-variable "CI_BASE_IMAGE_TAG_CONTENT" {
+# Supplemental tags only. ci-bake-rocm.sh leaves these empty when the same ref
+# is already the primary CI_BASE_IMAGE_TAG.
+variable "CI_BASE_IMAGE_TAG_COMMIT_EXTRA" {
+  default = ""
+}
+
+variable "CI_BASE_IMAGE_TAG_CONTENT_EXTRA" {
  default = ""
 }

@@ -357,7 +365,8 @@ target "ci-base-rocm-ci" {
  cache-from = concat(
    compact([
      CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "",
-      CI_BASE_IMAGE_TAG_CONTENT != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT}" : "",
+      CI_BASE_IMAGE_TAG_COMMIT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_COMMIT_EXTRA}" : "",
+      CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
      CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
    ]),
    # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
@@ -365,7 +374,7 @@ target "ci-base-rocm-ci" {
    get_cache_from_rocm_deps(),
  )
  cache-to = ["type=inline"]
-  tags     = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE])
+  tags     = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_COMMIT_EXTRA, CI_BASE_IMAGE_TAG_CONTENT_EXTRA, CI_BASE_IMAGE_TAG_STABLE])
  output   = ["type=registry"]
 }

@@ -80,6 +80,9 @@ vllm chat --url http://{vllm-serve-host}:{vllm-serve-port}/v1

 # Quick chat with a single prompt
 vllm chat --quick "hi"
+
+# Print TTFT and throughput statistics after each response
+vllm chat --stats
 ```

 See [vllm chat](./chat.md) for the full reference of all available arguments.
@@ -97,6 +100,9 @@ vllm complete --url http://{vllm-serve-host}:{vllm-serve-port}/v1

 # Quick complete with a single prompt
 vllm complete --quick "The future of AI is"
+
+# Print TTFT and throughput statistics after each response
+vllm complete --stats
 ```

 See [vllm complete](./complete.md) for the full reference of all available arguments.
@@ -297,7 +297,7 @@ The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,
 vLLM raises a clear `ImportError` at tokenizer load. The override applies to
 any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,
 `deepseek_v32`, `deepseek_v4`, …). Models that don't use the HF
-fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.
+fast tokenizer (`mistral`, `kimi_audio`) ignore the flag.

 Tokenizer-bound workloads — long shared prefixes, bursty short prompts,
 batch detokenization — see the largest wins. If your bottleneck is GPU
@@ -178,7 +178,7 @@ incoming requests. The `AsyncLLMEngine` is designed for online serving, where it
 can handle multiple concurrent requests and stream outputs to clients.

 The OpenAI-compatible API server uses the `AsyncLLMEngine`. There is also a demo
-API server that serves as a simpler example in [vllm/entrypoints/api_server.py](../../vllm/entrypoints/api_server.py).
+API server that serves as a simpler example in [examples/applications/api_server/server.py](../../examples/applications/api_server/server.py).

 The code for `AsyncLLMEngine` can be found in [vllm/engine/async_llm_engine.py](../../vllm/engine/async_llm_engine.py).

@@ -159,7 +159,7 @@ Priority is **1 = highest** (tried first).

 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
-| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
@@ -168,9 +168,9 @@ Priority is **1 = highest** (tried first).
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
-| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

@@ -226,6 +226,7 @@ MLA decode backends are selected using the standard
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
+| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -127,6 +127,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 | Architecture | Models | CG for Image | CG for Video | Dual-Path Graph |
 | ------------ | ------ | ------------ | ------------ | --------------- |
 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
+| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
 | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
 | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
@@ -134,8 +135,9 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
-| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ |
-| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ |
+| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
+| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
+| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ |

 !!! note
    Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
@@ -159,7 +159,7 @@ For simplicity, we assume `block_size=1` in this section.

 ### High level idea

-The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to catch the full blocks. That means the same tokens of different groups are cached and evicted independently.
+The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to cache the full blocks. That means the same tokens of different groups are cached and evicted independently.

 When a new request comes in, we check the cache hit prefix of each group, and return the intersection of these groups as the cached prefix of the request. See below for the detailed algorithm for checking the cache hit of one group & performing the intersection.

@@ -89,6 +89,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
+| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.experts.hpc.HPCExperts] |
 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
@@ -52,7 +52,7 @@ __device__ void paged_attention_kernel(
 )
 ```

-There are also a list of template arguments above the function
+There is also a list of template arguments above the function
 signature that are determined during compilation time. `scalar_t`
 represents the data type of the query, key, and value data elements,
 such as FP16. `HEAD_SIZE` indicates the number of elements in each
@@ -178,7 +178,7 @@ const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride
                    + physical_block_offset * x;
 ```

-Unlike to `q_ptr`, `k_ptr` in each thread will point to different
+Unlike `q_ptr`, `k_ptr` in each thread will point to different
 key token at different iterations. As shown above, that `k_ptr`
 points to key token data based on `k_cache` at assigned block,
 assigned head and assigned token.
--- a/Show More
+++ b/Show More