[CI/Build] Fix topk histogram build on SM75 (#46550 )

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> (cherry picked from commit 191826ec61)
feat: support to OpenMOSS-Team (#44124 )
2026-06-24 01:01:53 -07:00 · 2026-06-24 01:01:53 -07:00 · 2026-06-24 01:01:53 -07:00 · 2026-06-24 01:01:53 -07:00 · 2026-06-24 01:01:53 -07:00 · 2026-06-24 01:01:48 -07:00
1053 changed files with 23938 additions and 64154 deletions
@@ -2,16 +2,17 @@ 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,7 +6,6 @@ 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
@@ -27,7 +26,6 @@ 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 50m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
      pytest -x -v -s tests/models/language/generation -m cpu_model
      pytest -x -v -s tests/models/language/pooling -m cpu_model"

@@ -68,15 +68,13 @@ 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/e2e/test_cpu_linear_attn_chunked_prefix.py"
+      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"

 - label: CPU-Quantization Model Tests
  depends_on: []
@@ -91,13 +89,11 @@ 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_w8a8.py"
+      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
      
 - label: CPU-Distributed Tests (PP+TP)
  depends_on: []
@@ -1,25 +0,0 @@
-group: Zen CPU
-depends_on: []
-steps:
- label: Zen-CPU-Kernel Tests
-  depends_on: []
-  soft_fail: false
-  device: zen5
-  no_plugin: true
-  source_file_dependencies:
-  - setup.py
-  - docker/Dockerfile.cpu
-  - requirements/cpu.txt
-  - requirements/build/cpu.txt
-  - csrc/cpu/
-  - cmake/cpu_extension.cmake
-  - CMakeLists.txt
-  - vllm/model_executor/layers/utils.py
-  - vllm/platforms/cpu.py
-  - vllm/platforms/zen_cpu.py
-  - vllm/platforms/__init__.py
-  - tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-zen-cpu-test.sh 20m "
-      pytest -x -v -s tests/model_executor/test_cpu_unquantized_gemm_dispatch.py"
@@ -23,5 +23,4 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      pytest -v -s basic_correctness/test_cpu_offload.py &&
      pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
@@ -128,10 +128,10 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
      pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
      pytest -v -s lora/test_transformers_model.py &&
      pytest -v -s lora/test_chatglm3_tp.py &&
-      pytest -v -s lora/test_llama_tp.py::test_llama_lora &&
      pytest -s -v lora/test_minicpmv_tp.py'

 - label: LoRA Multimodal
@@ -1,27 +0,0 @@
-group: Models - Distributed
-depends_on: 
-  - image-build-xpu
-steps:
- label: Distributed Model Tests (2 GPUs)
-  key: distributed-model-tests-2-gpus
-  timeout_in_minutes: 50
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 2+
-    mem: 24+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/model_executor/model_loader/sharded_state_loader.py
-  - vllm/model_executor/models/
-  - tests/model_executor/model_loader/test_sharded_state_loader.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m "not slow_test"'
@@ -125,5 +125,7 @@ steps:
      pip install open-clip-torch --no-deps &&
      cd tests &&
      pytest -v -s models/multimodal/processing/test_tensor_schema.py
+      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
+      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
      --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
  parallelism: 4
@@ -15,10 +15,9 @@ 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 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_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_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=""
@@ -526,22 +525,6 @@ 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
@@ -598,7 +581,6 @@ 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

@@ -653,10 +635,6 @@ 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
@@ -698,14 +676,12 @@ 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}"
@@ -715,33 +691,19 @@ configure_ci_base_image_refs() {
    fi
    CI_BASE_IMAGE_TAG="${primary_tag}"
    if [[ "${primary_tag}" == "${content_tag}" ]]; then
-        CI_BASE_IMAGE_TAG_CONTENT_EXTRA=""
+        CI_BASE_IMAGE_TAG_CONTENT=""
    else
-        CI_BASE_IMAGE_TAG_CONTENT_EXTRA="${content_tag}"
+        CI_BASE_IMAGE_TAG_CONTENT="${content_tag}"
    fi
-    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
+    export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT 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 "${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
+        if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then
+            echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}"
        fi
        echo "ci_base content image tag: ${content_tag}"
        if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then
@@ -766,8 +728,8 @@ ci_base_candidate_refs() {
    printf '%s\n' \
        "${IMAGE_TAG:-}" \
        "${CI_BASE_IMAGE_TAG:-}" \
-        "${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}" \
-        "${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}" \
+        "${CI_BASE_IMAGE_TAG_COMMIT:-}" \
+        "${CI_BASE_IMAGE_TAG_CONTENT:-}" \
        "${CI_BASE_IMAGE_TAG_STABLE:-}" \
        | awk 'NF && !seen[$0]++'
 }
@@ -781,10 +743,6 @@ 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
@@ -859,10 +817,6 @@ 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
@@ -1044,104 +998,12 @@ 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
@@ -1157,23 +1019,16 @@ 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" = ""
-EOF
-        write_ci_base_metadata_labels "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
-        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
+    "vllm.ci_base.content_hash" = "${CI_BASE_CONTENT_HASH}"
  }
 }

@@ -1181,7 +1036,7 @@ EOF
    done

    BAKE_FILES+=(-f "${CI_BASE_LABEL_OVERRIDE_PATH}")
-    echo "Appended ci_base metadata label override for targets: ${ci_base_targets[*]}"
+    echo "Appended ci_base content-hash label override for targets: ${ci_base_targets[*]}"
 }

 uses_rocm_csrc_cache() {
@@ -1264,18 +1119,6 @@ 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
@@ -1698,13 +1541,7 @@ 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
-            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
+            return 0
        fi

        echo "Remote image exists but does not have the expected ci_base content hash."
@@ -1887,16 +1724,15 @@ 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,20 +367,6 @@ 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
@@ -440,24 +426,6 @@ 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}}"
@@ -557,7 +525,6 @@ 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=32
+export CMAKE_BUILD_PARALLEL_LEVEL=16

 # 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_cpu_fused_moe.py
+    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
    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,49 +7,10 @@ set -euox pipefail
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-48-95}
 NUMA_NODE=${NUMA_NODE:-1}
-AGENT_SLOT=${AGENT_SLOT:-}
-IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}"
+IMAGE_NAME="cpu-test-$NUMA_NODE"
 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; 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"' \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
    >/dev/null
 } 9>/tmp/docker-pull.lock

@@ -1,20 +0,0 @@
-#!/bin/bash
-
-# This script build the CPU docker image and run the offline inference inside the container.
-# It serves a sanity check for compilation and basic model usage.
-set -euox pipefail
-
-# allow to bind to different cores
-CORE_RANGE=${CORE_RANGE:-48-95}
-NUMA_NODE=${NUMA_NODE:-1}
-IMAGE_NAME="zen-cpu-test-$NUMA_NODE"
-TIMEOUT_VAL=$1
-TEST_COMMAND=$2
-
-# building the docker image
-echo "--- :docker: Building Docker image"
-docker build --progress plain --tag "$IMAGE_NAME" --target vllm-zen-test -f docker/Dockerfile.cpu .
-
-# Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
-        timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
@@ -90,16 +90,6 @@ 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
@@ -152,7 +142,6 @@ PY

 run_style_clippy() {
  install_cargo_sort
-  install_cargo_deny

  log_section "Checking Rust formatting"
  cargo fmt --manifest-path rust/Cargo.toml --all -- --check
@@ -160,13 +149,6 @@ 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,14 +33,6 @@ 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,6 +112,33 @@ 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
@@ -152,8 +179,72 @@ 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]
@@ -263,6 +354,54 @@ 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
@@ -390,6 +529,61 @@ 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
@@ -423,7 +617,7 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - pytest -v -s reasoning
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
  - pytest -v -s parser
  - pytest -v -s transformers_utils
@@ -438,7 +632,7 @@ steps:
 #-----------------------------------------------------  mi300 · basic_correctness  -----------------------------------------------------#

 - label: Basic Correctness # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
@@ -452,11 +646,11 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s basic_correctness/test_mem.py
-  - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py
+  - 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: 65
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -474,7 +668,7 @@ steps:
  - tests/model_executor/model_loader/test_sharded_state_loader.py
  - tests/models/
  commands:
-  - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+  - 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)'
@@ -488,7 +682,6 @@ 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/
@@ -601,7 +794,7 @@ steps:

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

- label: Platform Tests # TBD
+- label: Platform Tests (CUDA) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
@@ -632,27 +825,10 @@ 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
@@ -677,8 +853,8 @@ steps:
  - pytest -v -s distributed/test_eplb_execute.py
  - pytest -v -s distributed/test_eplb_spec_decode.py

- label: Distributed Tests (2xH100-2xMI300) # TBD
-  timeout_in_minutes: 30
+- label: Distributed Tests (2xH100-2xMI250) # TBD
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -689,19 +865,13 @@ 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
@@ -715,7 +885,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=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+  - TARGET_TEST_SUITE=A100 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
@@ -845,12 +1015,10 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/serve
-  - tests/entrypoints/scale_out
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  - pytest -v -s entrypoints/scale_out

 - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
  timeout_in_minutes: 180
@@ -1198,31 +1366,10 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt

- label: ROCm LM Eval Large Models (8 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_8
-  optional: true
-  num_gpus: 8
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/layers/quantization/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/model_executor/layers/layernorm.py
-  - csrc/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
-
 #---------------------------------------------------------  mi300 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
@@ -1235,7 +1382,7 @@ steps:
  commands:
    - pip install tensorizer
    # Basic
-    - python3 basic/offline_inference/chat.py
+    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
    - 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
@@ -1439,10 +1586,10 @@ steps:
  - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset

 - label: Basic Models Tests (Other) # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1484,7 +1631,6 @@ 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"
@@ -1689,54 +1835,6 @@ 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
@@ -2099,7 +2197,6 @@ 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:
@@ -2118,11 +2215,45 @@ 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: 20
+  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:
@@ -2148,28 +2279,11 @@ 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
@@ -2208,7 +2322,7 @@ steps:
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
  - tests/v1/kv_connector/nixl_integration/
  - vllm/platforms/rocm.py
  commands:
@@ -2220,10 +2334,9 @@ 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/
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
  - tests/v1/kv_connector/nixl_integration/
  - vllm/platforms/rocm.py
  commands:
@@ -2235,7 +2348,6 @@ 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
@@ -2272,7 +2384,7 @@ steps:
 #------------------------------------------------------  mi300 · weight_loading  -------------------------------------------------------#

 - label: Weight Loading Multiple GPU # TBD
-  timeout_in_minutes: 30
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -2284,7 +2396,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: 30
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_2
  num_gpus: 2
@@ -2413,6 +2525,27 @@ steps:
  - export VLLM_USE_DEEP_GEMM=0
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4

+- label: ROCm LM Eval Large Models (8 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
+  agent_pool: mi325_8
+  optional: true
+  num_gpus: 8
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/model_executor/layers/quantization/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/model_executor/layers/layernorm.py
+  - csrc/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
+
 #-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#

 - label: Language Models Test (Extended Generation) # TBD
@@ -2561,12 +2694,10 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/serve
-  - tests/entrypoints/scale_out
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  - pytest -v -s entrypoints/scale_out

 - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
  timeout_in_minutes: 180
@@ -2771,7 +2902,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/examples"
@@ -2783,7 +2914,7 @@ steps:
  commands:
  - pip install tensorizer
  # Basic
-  - python3 basic/offline_inference/chat.py
+  - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
  - 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
@@ -2828,7 +2959,6 @@ steps:
  - rocm-smi
  - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
  - pytest -v -s tests/kernels/attention/test_attention_selector.py
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py

 - label: Kernels Attention Test %N # TBD
  timeout_in_minutes: 60
@@ -3178,7 +3308,7 @@ steps:
 #------------------------------------------------------  mi355 · weight_loading  -------------------------------------------------------#

 - label: Weight Loading Multiple GPU # TBD
-  timeout_in_minutes: 30
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_2
  num_gpus: 2
@@ -3190,7 +3320,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: 30
+  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_2
  working_dir: "/vllm-workspace/tests"
@@ -16,9 +16,3 @@ 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,11 +11,6 @@ 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
-  key: platform-tests
+- label: Platform Tests (CUDA)
+  key: platform-tests-cuda
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
@@ -13,20 +13,6 @@ 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
@@ -50,19 +36,6 @@ 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
@@ -75,19 +48,6 @@ 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
@@ -100,19 +60,6 @@ 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
@@ -156,20 +103,6 @@ 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,21 +37,6 @@ 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
@@ -174,8 +159,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 (8xH100)
-  key: distributed-tests-8xh100
+- label: Distributed Tests (8 GPUs)(H100)
+  key: distributed-tests-8-gpus-h100
  timeout_in_minutes: 10
  device: h100
  num_devices: 8
@@ -195,8 +180,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 (4xA100)
-  key: distributed-tests-4xa100
+- label: Distributed Tests (4 GPUs)(A100)
+  key: distributed-tests-4-gpus-a100
  device: a100
  optional: true
  num_devices: 4
@@ -210,8 +195,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 (2xH100-2xMI300)
-  key: distributed-tests-2xh100-2xmi300
+- label: Distributed Tests (2 GPUs)(H100)
+  key: distributed-tests-2-gpus-h100
  timeout_in_minutes: 15
  device: h100
  optional: true
@@ -225,8 +210,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 (2xB200)
-  key: distributed-tests-2xb200
+- label: Distributed Tests (2 GPUs)(B200)
+  key: distributed-tests-2-gpus-b200
  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 (4xH100)
-  key: deepseek-v2-lite-sync-eplb-accuracy-4xh100
+- label: DeepSeek V2-Lite Sync EPLB Accuracy
+  key: deepseek-v2-lite-sync-eplb-accuracy
  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 (4xH100)
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
  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 (2xB200)
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
  timeout_in_minutes: 60
  device: b200-k8s
  optional: true
@@ -112,11 +112,6 @@ 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
@@ -29,8 +29,6 @@ steps:
  mirror:
    amd:
      device: mi325_1
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
      depends_on:
      - image-build-amd

@@ -42,12 +40,10 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/serve
-  - tests/entrypoints/scale_out
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  - pytest -v -s entrypoints/scale_out
  mirror:
    amd:
      device: mi325_1
@@ -14,16 +14,6 @@ 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
@@ -274,8 +274,8 @@ steps:
    - pytest -v -s kernels/helion/

 
- label: Kernels FP8 MoE Test (1xH100)
-  key: kernels-fp8-moe-test-1xh100
+- label: Kernels FP8 MoE Test (1 H100)
+  key: kernels-fp8-moe-test-1-h100
  timeout_in_minutes: 90
  device: h100
  num_devices: 1
@@ -291,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 (2xH100)
-  key: kernels-fp8-moe-test-2xh100
+- label: Kernels FP8 MoE Test (2 H100s)
+  key: kernels-fp8-moe-test-2-h100s
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
@@ -28,8 +28,7 @@ steps:
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py

-# - label: LM Eval Large Models (4xA100)
-#   key: lm-eval-large-models-4xa100
+# - label: LM Eval Large Models (4 GPUs)(A100)
 #   device: a100
 #   optional: true
 #   num_devices: 4
@@ -41,8 +40,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 (4xH100)
-  key: lm-eval-large-models-4xh100
+- label: LM Eval Large Models (4 GPUs)(H100)
+  key: lm-eval-large-models-4-gpus-h100
  device: h100
  optional: true
  num_devices: 4
@@ -54,8 +53,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 (2xB200)
-  key: lm-eval-small-models-2xb200
+- label: LM Eval Small Models (B200)
+  key: lm-eval-small-models-b200
  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
@@ -65,20 +64,8 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

- 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
+- label: LM Eval Large Models (B200, EP)
+  key: lm-eval-large-models-b200-ep
  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
@@ -89,8 +76,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 (2xB200)
-  key: lm-eval-qwen3-5-models-2xb200
+- label: LM Eval Qwen3.5 Models (B200)
+  key: lm-eval-qwen3-5-models-b200
  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
@@ -106,8 +93,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 (8xH200)
-  key: lm-eval-large-models-8xh200
+- label: LM Eval Large Models (H200)
+  key: lm-eval-large-models-h200
  timeout_in_minutes: 60
  device: h200
  optional: true
@@ -205,8 +192,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) (2xH100)
-  key: gpqa-eval-gpt-oss-2xh100
+- label: GPQA Eval (GPT-OSS) (H100)
+  key: gpqa-eval-gpt-oss-h100
  timeout_in_minutes: 120
  device: h100
  optional: true
@@ -219,8 +206,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) (2xB200)
-  key: gpqa-eval-gpt-oss-2xb200
+- label: GPQA Eval (GPT-OSS) (B200)
+  key: gpqa-eval-gpt-oss-b200
  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
@@ -239,8 +226,6 @@ steps:
  device: dgx-spark
  optional: true
  num_devices: 1
-  depends_on:
-  - arm64-image-build
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -224,16 +224,6 @@ 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
@@ -260,12 +250,6 @@ 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
@@ -278,16 +262,6 @@ 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
@@ -370,7 +344,7 @@ steps:
  - pytest -v -s test_ray_env.py
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
-  - pytest -v -s reasoning
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
  - pytest -v -s tokenizers_
  - pytest -v -s parser
@@ -23,16 +23,3 @@ 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,11 +45,6 @@ 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,10 +15,6 @@ 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
@@ -36,21 +32,6 @@ 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,17 +109,6 @@ 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,11 +195,3 @@ 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 -k "not test_invalid_json_schema and not test_invalid_regex"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"

  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
@@ -46,7 +46,7 @@ steps:
  - vllm/v1/engine/
  - tests/utils.py
  # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
-  - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+  - tests/entrypoints/serve/disagg/test_serving_tokens.py
  - tests/entrypoints/serve/instrumentator/test_basic.py
  - tests/entrypoints/serve/instrumentator/test_metrics.py
  # - tests/entrypoints/serve/dev/test_sleep.py
@@ -55,7 +55,7 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
  - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
-  - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
+  - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
  - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
  # - pytest -v -s entrypoints/serve/dev/test_sleep.py

@@ -94,8 +94,6 @@ steps:
    amd:
      device: mi325_1
      timeout_in_minutes: 65
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -13,13 +13,6 @@ 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 @ivanium
+/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /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 @ivanium
+/vllm/config/cache.py @heheda12345

 # Config utils
 /vllm/config/utils.py @hmellor
@@ -67,17 +67,16 @@
 /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 @ivanium
+/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
 /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 @ivanium
+/vllm/v1/kv_cache_interface.py @heheda12345
 /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 @ivanium
+/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche

 # Model runner V2
 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256
@@ -104,14 +103,13 @@
 /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 @ivanium
+/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
 /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 @ivanium
+/tests/v1/kv_connector @ApostaC @orozery
 /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,5 +3,3 @@
 self-hosted-runner:
  labels:
    - vllm-runners
-    # Not yet in actionlint's known-label set.
-    - macos-26
@@ -327,7 +327,7 @@ jobs:
                message: 'CC {users} for ROCm-related issue',
              },
              mistral: {
-                users: ['patrickvonplaten', 'juliendenize', 'andylolu2', 'NickLucche'],
+                users: ['patrickvonplaten', 'juliendenize', 'andylolu2'],
                message: 'CC {users} for Mistral-related issue',
              },
              // Add more label -> user mappings here
@@ -11,25 +11,13 @@ permissions:

 jobs:
  macos-m1-smoke-test:
-    # 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 }}
+    runs-on: macos-latest
    timeout-minutes: 30

    steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@v6.0.1

-      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      - uses: astral-sh/setup-uv@v7
        with:
          enable-cache: true
          cache-dependency-glob: |
@@ -84,11 +72,14 @@ jobs:
          # Test health endpoint
          curl -f http://localhost:8000/health

-          # 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 \
+          # Test completion
+          curl -f http://localhost:8000/v1/completions \
            -H "Content-Type: application/json" \
-            -d "$PAYLOAD"
+            -d '{
+              "model": "Qwen/Qwen3-0.6B",
+              "prompt": "Hello",
+              "max_tokens": 5
+            }'

          # Cleanup
          kill "$SERVER_PID"
@@ -48,7 +48,7 @@ jobs:
    if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
    runs-on: [self-hosted, linux, x64, vllm-runners]
    steps:
-    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
      with:
        python-version: "3.12"
@@ -131,19 +131,6 @@ repos:
        --python-version, "3.12",
      ]
      files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$
-    - id: pip-compile
-      alias: pip-compile-cpu
-      name: pip-compile-cpu
-      args: [
-        requirements/test/cuda.in,
-        -o, requirements/test/cpu.txt,
-        --index-strategy, unsafe-best-match,
-        --torch-backend, cpu,
-        --python-platform, x86_64-manylinux_2_28,
-        --python-version, "3.12",
-      ]
-      files: ^requirements/(common|cpu|test/(cuda|cpu))\.(in|txt)$
-      exclude: ^requirements/test/cuda\.txt$
    - id: pip-compile
      alias: pip-compile-docs
      name: pip-compile-docs
@@ -140,21 +140,6 @@ if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
    WITH_SOABI)
 endif()

-#
-# fs_io extension (pure CXX; must stay above the non-CUDA device branch
-# so CPU builds define the target before the early return).
-# GIL-releasing filesystem helpers for FileSystemTierManager.
-#
-if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
-  define_extension_target(
-    fs_io_C
-    DESTINATION vllm
-    LANGUAGE CXX
-    SOURCES csrc/fs_io.cpp
-    USE_SABI 3.11
-    WITH_SOABI)
-endif()
-
 #
 # Forward the non-CUDA device extensions to external CMake scripts.
 #
@@ -285,16 +270,6 @@ 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()

 #
@@ -894,9 +869,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;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}")
  else()
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;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")
@@ -80,17 +80,13 @@ _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. 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
-]
+# Backends to benchmark
+FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]


 def setup_flashinfer_workspace(
@@ -999,10 +995,7 @@ def main():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

-    # 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}")
+    device = torch.device(f"cuda:{rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)

@@ -391,19 +391,16 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int
        config = dict(zip(keys, config_values))
        configs.append(config)

-    # 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.
+    # 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
    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[:]:
-            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):
+            if (
+                config["BLOCK_SIZE_K"] % block_k != 0
+                or config["BLOCK_SIZE_N"] % block_n != 0
+            ):
                configs.remove(config)
    return configs

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

@@ -15,15 +14,17 @@ 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 = ["vec"]
-if torch.cpu._is_amx_tile_supported():
-    ISA_CHOICES.append("amx")
-if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
-    ISA_CHOICES.append("neon")
+ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]


@torch.inference_mode()
@@ -144,7 +145,7 @@ if __name__ == "__main__":
        "--isa",
        type=str,
        choices=ISA_CHOICES,
-        default="vec",
+        default=ISA_CHOICES[0],
        help=f"ISA to use (available: {ISA_CHOICES})",
    )
    parser.add_argument("--seed", type=int, default=0)
@@ -24,10 +24,7 @@ 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
@@ -169,13 +166,12 @@ elseif (S390_FOUND)
        "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
    message(STATUS "RISC-V detected")
-    if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0)
+    if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
        message(FATAL_ERROR
-            "VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'")
+            "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
    endif()
    # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
-    # 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.
+    # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
    if(NOT DEFINED VLLM_RVV_VLEN)
        # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
        if(EXISTS /proc/cpuinfo)
@@ -427,7 +423,6 @@ 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()

@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
+          GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
          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(cpu_utils::get_max_threads()),
+        thread_num(omp_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 = cpu_utils::get_max_threads();
+    const int32_t thread_num = omp_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,10 +887,12 @@ 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,
@@ -1423,7 +1425,7 @@ class AttentionMainLoop {

 public:
  void operator()(const AttentionInput* input) {
-    const int thread_num = cpu_utils::get_max_threads();
+    const int thread_num = omp_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,16 +50,7 @@ FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
  b1 = (__vector float)vec_mergel(zeros, raw);
 }

-// [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];
-}
+// Note: c10::Half (FP16) is not supported on PowerPC architecture

 template <int32_t M, typename kv_cache_t>
 FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
@@ -1,3 +1,5 @@
+#include <sleef.h>
+
 #include "cpu/cpu_types.hpp"
 #include "cpu/utils.hpp"
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -14,18 +16,6 @@
  #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) {                                               \
@@ -35,7 +25,6 @@
            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.");                  \
      }                                                               \
@@ -70,12 +59,10 @@ 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);
@@ -85,15 +72,8 @@ 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));
@@ -194,7 +174,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
-      auto tanh_vec = inner_vec.tanh();
+      vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
      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);
@@ -260,14 +240,13 @@ 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 = cpu_utils::get_max_threads();
+  const int32_t thread_num = omp_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));
@@ -287,18 +266,12 @@ 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 + optional packed input + weight
+    // input tile + weight
    const int32_t n_size_cache_limit =
-        (cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) /
-        (input_size_2 * sizeof(scalar_t));
+        (cache_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>(
@@ -351,9 +324,6 @@ 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));
@@ -433,8 +403,7 @@ 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 =
-          gemm_t::WeightOCGroupSize * input_size_13;
+      const int32_t w13_n_group_stride = 16 * input_size_13;
      const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;

      for (;;) {
@@ -497,23 +466,8 @@ 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_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;
+          // copy inputs
+          {
            scalar_t* __restrict__ curr_w13_input_buffer_iter =
                curr_w13_input_buffer;
            for (int32_t i = 0; i < actual_token_num; ++i) {
@@ -545,12 +499,14 @@ 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_gemm_input_buffer, w13_weight_ptr_0_iter,
+              gemm.gemm(curr_w13_input_buffer_iter, 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);
@@ -563,7 +519,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_gemm_input_buffer, w13_weight_ptr_1_iter,
+              gemm.gemm(curr_w13_input_buffer_iter, 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);
@@ -616,8 +572,7 @@ 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 =
-          gemm_t::WeightOCGroupSize * input_size_2;
+      const int32_t w2_n_group_stride = 16 * input_size_2;

      for (;;) {
        int32_t task_id = counter_ptr->acquire_counter();
@@ -656,30 +611,13 @@ 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_w2_gemm_input_buffer, w2_weight_ptr_iter,
+            gemm.gemm(curr_w13_gemm_output_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,20 +25,4 @@
  #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,26 +497,6 @@ 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];
@@ -535,15 +515,6 @@ 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];
@@ -962,4 +933,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,17 +3,13 @@

 // VLEN-to-LMUL mapping for RISC-V Vector extension.
 //
-// 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
+// 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

 #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
@@ -21,7 +17,6 @@
  #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
@@ -46,16 +41,6 @@

 // ---- 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,13 +363,6 @@ 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 {};

-#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__)
+// 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_FLOATING_TYPES(TYPE, NAME, ...) \
  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -34,87 +34,6 @@ 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>{}), ...);
@@ -170,19 +89,6 @@ 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;
@@ -218,11 +124,13 @@ 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 {
@@ -234,16 +142,20 @@ 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);
    }
@@ -326,15 +238,6 @@ 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;
@@ -507,9 +410,8 @@ 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) {
@@ -533,6 +435,7 @@ 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]);
@@ -599,20 +502,28 @@ 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));

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

+    // 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],
@@ -715,16 +626,6 @@ 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;
@@ -758,7 +659,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
 };

 struct INT8Vec16 : public Vec<INT8Vec16> {
-  constexpr static int VEC_NUM_ELEM = 16;
+  constexpr static int VEC_NUM_ELEM = 16;  // 128 bits / 8 bits = 16

  union AliasReg {
    __vector signed char reg;
@@ -806,11 +707,6 @@ 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;
@@ -827,15 +723,6 @@ 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
@@ -882,39 +769,38 @@ 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) {
-  __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);
+  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);
 }

 inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
-  __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);
+  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);
 }

 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
@@ -978,6 +864,7 @@ inline void prefetch(const void* addr) {

 struct INT8Vec64 {
  __vector signed char data[4];
+
  INT8Vec64() = default;

  explicit INT8Vec64(const int8_t* ptr) {
@@ -1013,4 +900,5 @@ struct INT8Vec64 {
  void nt_save(int8_t* ptr) const { save(ptr); }
 };
 }  // namespace vec_op
+
 #endif
@@ -3,7 +3,6 @@
 #define CPU_TYPES_X86_HPP

 #include <immintrin.h>
-#include <sleef.h>
 #include <torch/all.h>

 #ifndef __AVX2__
@@ -593,8 +592,6 @@ 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); }
@@ -792,12 +789,6 @@ 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 = cpu_utils::get_max_threads();
+  const int32_t thread_num = omp_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 = cpu_utils::get_max_threads();
+  const int64_t thread_num = omp_get_max_threads();
  if (num_tokens > thread_num) {
 #pragma omp parallel for
    for (int64_t i = 0; i < num_tokens; ++i) {
@@ -213,8 +213,6 @@ 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,9 +21,6 @@ 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) {
@@ -1,503 +0,0 @@
-#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,8 +104,6 @@ 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,9 +18,17 @@ 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__
@@ -251,7 +259,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 = cpu_utils::get_max_threads();
+  const int max_threads = omp_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,3 +1,4 @@
+
 #include "cpu_types.hpp"

 namespace {
@@ -96,91 +97,6 @@ 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
@@ -258,75 +174,6 @@ 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,18 +289,19 @@ 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,                                                         \
-      has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
-      has_initial_states_value,                                                                     \
-      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,                         \
+      nullptr,                                                      \
+      false,                                                        \
+      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,
@@ -342,9 +343,6 @@ 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);
@@ -375,7 +373,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
            width,
            dim,
            seqlen,
-            has_initial_state[bs]);
+            /* has_initial_state */ false);
      }
    });
  }
@@ -285,125 +285,6 @@ 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,
@@ -455,11 +336,6 @@ 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,19 +9,11 @@
 #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,89 +208,6 @@ 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,16 +237,6 @@ 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,
@@ -548,7 +538,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif

  // fused moe
-#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
+#if defined(__AVX512F__)
  ops.def(
      "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
      "-> ()");
@@ -559,7 +549,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  // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
+#endif
  ops.def(
      "mla_decode_kvcache("
      "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -609,19 +599,6 @@ 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,14 +2,13 @@
 #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, NEON };
+enum class ISA { AMX, VEC, RVV };

 inline ISA get_isa(const std::string& isa) {
  if (isa == "amx") {
@@ -18,8 +17,6 @@ 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, CUdeviceptr addr = 0) {
-  CUresult status = cuMemAddressReserve(d_mem, size, alignment, addr, 0);
+                                     size_t alignment) {
+  CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 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, addr, 0);
+  status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
  if (status != CUresult(0)) {
    return status;
  }
@@ -535,14 +535,7 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {
  Py_DECREF(py_result);
  PyGILState_Release(gstate);

-  // 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);
-  }
+  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.
@@ -555,13 +548,11 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {
  unmap_and_release(device, size, d_mem, p_memHandle);
 #endif

-  // Free the virtual address. On ROCm this also covers an asleep allocation,
-  // whose placeholder reservation made by sleep() is still held here.
+  // free address and the handle
  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]);
  }
@@ -681,29 +672,6 @@ 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
@@ -768,7 +736,6 @@ 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);

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

-  // 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) {
+  if (batch_fn != nullptr) {
    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
@@ -20,26 +20,20 @@ __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] or
-                                          // [num_groups, hidden_size];
-                                          // null if !HasWeight
-    const int64_t weight_stride,          // 0 or weight.stride(0)
+    const scalar_t* __restrict__ weight,  // [hidden_size], null if !HasWeight
    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);
@@ -48,7 +42,6 @@ __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) {
@@ -76,7 +69,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_row);
+  auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
  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;
@@ -218,24 +211,15 @@ 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, double epsilon) {
+              std::optional<torch::stable::Tensor> weight,  // [hidden_size]
+              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);
@@ -272,16 +256,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, weight_stride, epsilon,
-                      num_tokens, hidden_size);
+                      input_shape_d3, weight_ptr, 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, /*weight_stride=*/0, epsilon,
-                      num_tokens, hidden_size);
+                      input_shape_d3, weight_ptr, epsilon, num_tokens,
+                      hidden_size);
            }
          });
        });
@@ -11,7 +11,6 @@
 #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))
@@ -350,102 +349,19 @@ __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_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;
-
+__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;
 #pragma unroll
    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]);
+      x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
    }
-
-    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);
+    out[token_idx * d + idx] = x;
  }
 }

@@ -710,82 +626,52 @@ 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 int64_t num_tokens = output.numel() / hidden_size;
+  const auto 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());

-#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)
+  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;

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

-    // 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
+    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;
+  }
 }

 void moe_lora_align_block_size(
@@ -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 < 120) {
+  if (cuda_device_capability >= 100 && cuda_device_capability < 110) {
    return CUDA_VERSION >= 12080;
  }
  #endif
@@ -119,10 +119,6 @@ __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));
@@ -157,10 +153,6 @@ __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) {
@@ -217,56 +209,45 @@ 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);

-#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);     \
-      }                                                                 \
-    }                                                                   \
+#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);                                          \
+      }                                                                    \
+    }                                                                      \
  } while (0)

  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
@@ -281,7 +262,6 @@ 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.
@@ -326,12 +306,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))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #endif

  if (mn_idx >= tma_aligned_mn) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaTriggerProgrammaticLaunchCompletion();
+    asm volatile("griddepcontrol.launch_dependents;");
 #endif
    return;
  }
@@ -448,7 +428,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))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -41,6 +41,12 @@ 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,29 +283,6 @@ __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,30 +59,11 @@ 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 {
-  // 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)
+  F16 = 0,
+  INT8 = 1,
+  INT6 = 2,
+  INT4 = 3,
 };

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

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

 namespace quickreduce {

@@ -207,168 +206,6 @@ 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.
@@ -540,6 +377,7 @@ 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 = MIN(laneid / 4, HEAD_SIZE / 8 - 1);
+    const int qhead_elemh8 = laneid / 4;

    for (int h = 0; h < QHLOOP - 1; h++) {
      const int qhead_idx = h * 4 + lane4id;
@@ -70,15 +70,6 @@ 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>
@@ -1246,45 +1237,6 @@ 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());
@@ -1299,21 +1251,37 @@ 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:
-        WVSPLIT_TILE(sYT, 1)
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 1)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 1)
        break;
      case 2:
-        WVSPLIT_TILE(sYT, 2)
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 2)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 2)
        break;
      case 3:
-        WVSPLIT_TILE(sYT, 3)
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 3)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 3)
        break;
      case 4:
-        WVSPLIT_TILE(sYT, 4)
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 4)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 4)
        break;
      case 5:
-        WVSPLIT_TILE(sYT, 5)
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 5)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 5)
        break;
      default:
        throw std::runtime_error(
@@ -311,10 +311,9 @@ 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,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
    bash build_rust.sh
-
 #################### RUST BUILD IMAGE ####################

 #################### CSRC BUILD IMAGE ####################
@@ -793,7 +792,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.13
+ARG FLASHINFER_VERSION=0.6.12
 RUN --mount=type=cache,target=/opt/uv/cache \
    uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -899,6 +898,13 @@ 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).
@@ -193,16 +193,26 @@ FROM base AS vllm-test-deps

 WORKDIR /vllm-workspace

-# Test requirements are compiled from requirements/test/cuda.in into
-# requirements/test/cpu.txt by the pip-compile-cpu pre-commit hook, which
-# resolves CPU wheels via uv's --torch-backend cpu.
-COPY requirements/test/cpu.txt requirements/test/cpu.txt
+# Copy test requirements
+COPY requirements/test/cuda.in requirements/test/cpu.in

-# cpu.txt is compiled for x86_64, so platform markers are resolved away. Drop
-# packages unavailable on aarch64 (decord, terratorch) for arm builds.
-RUN case "$(uname -m)" in \
-    aarch64|arm64) sed -i '/^decord==/d; /^terratorch==/d' requirements/test/cpu.txt ;; \
-    esac
+RUN \
+    sed -i '/mamba_ssm/d' requirements/test/cpu.in && \
+    remove_packages_not_supported_on_aarch64() { \
+    case "$(uname -m)" in \
+    aarch64|arm64) \
+    sed -i '/decord/d' requirements/test/cpu.in; \
+    sed -i '/terratorch/d' requirements/test/cpu.in; \
+    ;; \
+    esac; \
+    }; \
+    remove_packages_not_supported_on_aarch64 && \
+    sed -i 's/^torch==.*/torch==2.11.0/g' requirements/test/cpu.in && \
+    sed -i 's/torchaudio.*/torchaudio/g' requirements/test/cpu.in && \
+    sed -i 's/torchvision.*/torchvision/g' requirements/test/cpu.in && \
+    # Related issue: https://github.com/vllm-project/vllm/pull/38800#issuecomment-4228314305
+    sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/cpu.in && \
+    uv pip compile requirements/test/cpu.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu

 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -r requirements/test/cpu.txt
@@ -318,30 +328,3 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install "vllm[zen]"

 ENTRYPOINT ["vllm", "serve"]
-
-######################### ZEN CPU TEST IMAGE #########################
-FROM vllm-openai-zen AS vllm-zen-test
-
-WORKDIR /vllm-workspace
-
-COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/test/cpu.txt
-
-ADD ./tests/ ./tests/
-ADD ./examples/ ./examples/
-ADD ./benchmarks/ ./benchmarks/
-ADD ./vllm/collect_env.py .
-ADD ./docker/ ./docker/
-ADD ./.buildkite/ ./.buildkite/
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -e tests/vllm_test_utils
-
-# enable fast downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE=1
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT=60
-
-ENTRYPOINT []
@@ -123,10 +123,9 @@ 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,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
    bash build_rust.sh
-
 #################### RUST BUILD IMAGE ####################

 #################### WHEEL BUILD IMAGE ####################
@@ -257,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.


 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.13
+# release version: v0.6.12
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=cache,target=/root/.cache/uv \
    echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.13 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
    && cd flashinfer \
    && git submodule update --init --recursive \
    && echo "finish git clone flashinfer..." \
@@ -575,9 +575,6 @@ 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 \
@@ -698,9 +695,6 @@ 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.16.post2"
+ARG AITER_BRANCH="v0.1.13.post1"
 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 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=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,26 +28,22 @@ 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,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
    bash build_rust.sh

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

 WORKDIR /workspace/

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

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

 # 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 && \
@@ -88,7 +82,6 @@ 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
@@ -108,13 +101,7 @@ 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 && \
-    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
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest

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



-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
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
 CMD ["/bin/bash"]

 ######################### UCX + NIXL BUILD STAGE #########################
@@ -217,7 +204,10 @@ 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
+    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}"

 # Keep source-dependent layers near the end so frequent code-only changes
 # don't invalidate heavy dependency and UCX/NIXL layers.
@@ -302,22 +302,14 @@ 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. 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.
+# 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.
 variable "CI_BASE_IMAGE_TAG" {
  default = "rocm/vllm-dev:ci_base"
 }

-# 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" {
+variable "CI_BASE_IMAGE_TAG_CONTENT" {
  default = ""
 }

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

@@ -68,7 +68,7 @@
      "default": "true"
    },
    "FLASHINFER_VERSION": {
-      "default": "0.6.13"
+      "default": "0.6.12"
    },
    "GDRCOPY_CUDA_VERSION": {
      "default": "12.8"
@@ -80,9 +80,6 @@ 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.
@@ -100,9 +97,6 @@ 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`, `kimi_audio`) ignore the flag.
+fast tokenizer (`mistral`, `grok2`, `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 [examples/applications/api_server/server.py](../../examples/applications/api_server/server.py).
+API server that serves as a simpler example in [vllm/entrypoints/api_server.py](../../vllm/entrypoints/api_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 |
@@ -167,11 +167,10 @@ Priority is **1 = highest** (tried first).
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %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`, `int4_per_token_head`, `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`, `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 |

@@ -221,13 +220,12 @@ MLA decode backends are selected using the standard
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
 | `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,7 +127,6 @@ 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` | ✅︎ | ❌︎ | ❌︎ |
@@ -135,9 +134,8 @@ 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`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
-| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
-| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ |
+| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ |
+| `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.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Mohammad Miadh Angkadandkhluu	7b3d595eb1	[CI/Build] Fix topk histogram build on SM75 (#46550 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> (cherry picked from commit `191826ec61`)	2026-06-24 01:01:53 -07:00
hurukawaandkhluu	e66b766bc4	feat: support to OpenMOSS-Team (#44124 ) Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa19 <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> (cherry picked from commit `489abadfb8`)	2026-06-24 01:01:53 -07:00
Jee Jee Liandkhluu	51000a95a1	[Kernel] GLM5 Router GEMM (#46385 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> (cherry picked from commit `9d6fdc2901`)	2026-06-24 01:01:53 -07:00
Roberto L. Castroandkhluu	2e0c5f52f7	[Perf][DSv4/DSv3.2] Add cluster-cooperative topK kernel for low-latency scenarios (#43008 ) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> (cherry picked from commit `855cd4d787`)	2026-06-24 01:01:53 -07:00
Yongye Zhuandkhluu	0085058ca8	[Kernel] Add FlashInferCutedslMxfp8LinearKernel (cute-dsl mm_mxfp8) (#46393 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `11b56b2ff2`)	2026-06-24 01:01:53 -07:00
Gabriel Wuandkhluu	fdd8e4efd3	[Bugfix] Allow flashinfer_cutlass as a clamped NVFP4 MoE backend (#46492 ) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> (cherry picked from commit `0d4d164488`) Signed-off-by: khluu <khluu000@gmail.com>	2026-06-24 01:01:48 -07:00
Yongye Zhuandkhluu	c6561c2fb6	[Minimax-M3] BF16/FP8 Indexer using MSA (#45892 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> (cherry picked from commit `6691f087a6`)	2026-06-24 00:53:06 -07:00