[Model] Add Inkling multi-depth MTP support [5/N]

Extend the merged Inkling MTP=1 implementation to multiple checkpoint depths with multi-step speculative decoding and KV-cache plumbing. Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-07-16 21:17:29 +00:00
1424 changed files with 19418 additions and 123689 deletions
@@ -1,103 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Audit vLLM compiled libraries for PyTorch stable ABI compliance."""
-
-import fnmatch
-import sys
-from pathlib import Path
-
-from torch_abi_audit import inspect_package
-from torch_abi_audit.report import ExtensionReport, PackageReport
-
-# Temporary allowlist of extensions not yet on the stable ABI.
-# Shrink and remove over time.
-ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
-    "_flashkda_C.abi3.so",
-    "vllm_flash_attn/_vllm_fa2_C.abi3.so",
-    "vllm_flash_attn/_vllm_fa3_C.abi3.so",
-    "third_party/deep_gemm/_C*.so",
-)
-
-
-def _relative_path(lib: ExtensionReport, package_root: Path) -> str:
-    try:
-        return lib.path.relative_to(package_root).as_posix()
-    except ValueError:
-        return lib.path.name
-
-
-def _is_torch_unstable(lib: ExtensionReport) -> bool:
-    return lib.error is None and lib.torch.uses_torch and not lib.torch.stable
-
-
-def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool:
-    return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns)
-
-
-def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]:
-    return (*report.extensions, *report.bundled_libs)
-
-
-def _collect_unstable(report: PackageReport) -> list[str]:
-    return sorted(
-        _relative_path(lib, report.root)
-        for lib in _iter_libs(report)
-        if _is_torch_unstable(lib)
-    )
-
-
-def _find_stale_allowlist_entries(
-    report: PackageReport, patterns: tuple[str, ...]
-) -> list[str]:
-    """Allowlist patterns that match a built library which is no longer unstable."""
-    stale: list[str] = []
-    for pattern in patterns:
-        for lib in _iter_libs(report):
-            if lib.error is not None:
-                continue
-            if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern):
-                continue
-            if not _is_torch_unstable(lib):
-                stale.append(pattern)
-                break
-    return stale
-
-
-def check_torch_abi(
-    package: str = "vllm",
-    patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES,
-) -> int:
-    report = inspect_package(package)
-    if report.error:
-        print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr)
-        return 2
-
-    unstable = _collect_unstable(report)
-    unexpected = [
-        rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)
-    ]
-    stale = _find_stale_allowlist_entries(report, patterns)
-
-    if unexpected or stale:
-        if unexpected:
-            print(
-                "Not allowed: torch-unstable libraries outside "
-                f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}",
-                file=sys.stderr,
-            )
-        if stale:
-            print(
-                "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: "
-                f"{', '.join(stale)}",
-                file=sys.stderr,
-            )
-        return 1
-
-    print("Torch stable ABI check passed.")
-    return 0
-
-
-if __name__ == "__main__":
-    print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance")
-    sys.exit(check_torch_abi())
@@ -14,7 +14,6 @@ run_all_patterns:
  - "setup.py"
  - "csrc/"
  - "cmake/"
-  - ".buildkite/check-torch-abi.py"
 run_all_exclude_patterns:
  - "docker/Dockerfile."
  - "csrc/cpu/"
@@ -18,8 +18,6 @@ steps:
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
  - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
  - tests/kernels/mamba/test_cpu_short_conv.py
-  - tests/kernels/mamba/test_causal_conv1d.py
-  - tests/kernels/mamba/test_mamba_ssm.py
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -30,9 +28,7 @@ steps:
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
-      pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"

 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
@@ -1,26 +0,0 @@
-group: Benchmarks
-depends_on: 
-  - image-build-xpu
-steps:
- label: Benchmarks CLI Test
-  key: benchmarks-cli-test
-  timeout_in_minutes: 40
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/benchmarks/
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s benchmarks/'
@@ -2,44 +2,6 @@ group: Engine Intel
 depends_on:
  - image-build-xpu
 steps:
- label: Engine
-  key: engine
-  timeout_in_minutes: 40
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/config/
-  - vllm/engine/
-  - vllm/entrypoints/logger.py
-  - vllm/envs.py
-  - vllm/logger.py
-  - vllm/logging_utils/
-  - vllm/platforms/
-  - vllm/sequence.py
-  - vllm/triton_utils/
-  - vllm/utils/
-  - tests/engine
-  - tests/test_sequence
-  - tests/test_config
-  - tests/test_logger
-  - tests/test_vllm_port
-  - tests/test_jit_monitor.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py'
-
 - label: Engine (1 GPU)
  timeout_in_minutes: 30
  device: intel_gpu
@@ -61,41 +23,3 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
-
- label: V1 e2e (2 GPUs)
-  timeout_in_minutes: 30
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 2+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/compilation/
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/engine/
-    - vllm/envs.py
-    - vllm/forward_context.py
-    - vllm/inputs/
-    - vllm/logger.py
-    - vllm/logging_utils/
-    - vllm/model_executor/
-    - vllm/multimodal/
-    - vllm/platforms/
-    - vllm/sampling_params.py
-    - vllm/transformers_utils/
-    - vllm/triton_utils/
-    - vllm/utils/
-    - vllm/v1/
-    - tests/v1/e2e/spec_decode
-  commands:
-  - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
@@ -125,13 +125,13 @@ steps:
      pytest -v -s v1/kv_offload &&
      pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'

- label: NixlConnector PD accuracy (4 GPUs)
+- label: NixlConnector PD accuracy (2 GPUs)
  timeout_in_minutes: 60
-  num_devices: 4
+  num_devices: 2
  device: intel_gpu
  agent_tags:
    label: production
-    gpu: 4+
+    gpu: 2+
    mem: 16+
  no_plugin: true
  working_dir: "."
@@ -148,10 +148,7 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
-      PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
-      PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
-      PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
+      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'

 - label: Regression
  key: regression
@@ -262,25 +259,3 @@ steps:
    pytest -v -s detokenizer &&
    pytest -v -s -m "not cpu_test" ./multimodal &&
    pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
-
- label: Fusion Unit Tests
-  timeout_in_minutes: 30
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/compilation/
-    - tests/compile/passes/test_qk_norm_rope_fusion.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
@@ -1,33 +0,0 @@
-group: Model Executor Intel
-depends_on:
-  - image-build-xpu
-steps:
- label: Model Executor (Intel)
-  key: model-executor-intel
-  timeout_in_minutes: 45
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    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/engine/arg_utils.py
-  - vllm/config/model.py
-  - vllm/model_executor
-  - tests/model_executor
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'apt-get update && apt-get install -y curl libsodium23 &&
-      pip3 install tensorizer==2.10.1 &&
-      pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      export PYTHONFAULTHANDLER=1 &&
-      cd tests &&
-      pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"'
@@ -8,7 +8,7 @@ steps:
  agent_tags:
    label: production
    gpu: 2+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -28,9 +28,7 @@ steps:
      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
      cd tests &&
      pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
-      pytest -v -s v1/e2e/general/test_context_length.py &&
      ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
-      pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
      pytest -v -s v1/e2e/general/test_min_tokens.py'

 - label: Model Runner V2 Examples (Intel)
@@ -62,55 +60,3 @@ steps:
      python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
      python3 generate/multimodal/vision_language_offline.py --seed 0 &&
      python3 features/automatic_prefix_caching/prefix_caching_offline.py'
-
- label: Model Runner V2 Distributed (2 GPUs)
-  timeout_in_minutes: 50
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 2+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/v1/worker/gpu/
-    - vllm/v1/worker/gpu_worker.py
-    - tests/basic_correctness/test_basic_correctness.py
-    - tests/v1/distributed/test_async_llm_dp.py
-    - tests/v1/distributed/test_eagle_dp.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd tests &&
-      TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
-
- label: Model Runner V2 Spec Decode
-  timeout_in_minutes: 50
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    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/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_rejection_sampler_utils.py
-  - tests/v1/e2e/spec_decode/test_spec_decode.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd tests &&
-      pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
@@ -1,29 +0,0 @@
-group: Samplers Intel
-depends_on:
-  - image-build-xpu
-steps:
- label: Samplers Test (FlashInfer)
-  key: samplers-test-flashinfer-intel
-  timeout_in_minutes: 40
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    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/layers
-  - vllm/sampling_metadata.py
-  - tests/samplers
-  - tests/conftest.py
-  - vllm/entrypoints/generate/beam_search
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers'
@@ -145,32 +145,7 @@ steps:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'cd tests &&
-        pytest -v -s quantization/test_auto_round.py &&
-        pytest -v -s quantization/test_online.py'
-  - label: "XPU GPQA Eval (GPT-OSS)"
-    depends_on:
-      - image-build-xpu
-    timeout_in_minutes: 60
-    device: intel_gpu
-    agent_tags:
-      label: production
-      gpu: 1+
-      mem: 24+
-    no_plugin: true
-    env:
-      REGISTRY: "public.ecr.aws/q9t5s3a7"
-      REPO: "vllm-ci-test-repo"
-      VLLM_TEST_DEVICE: "xpu"
-    source_file_dependencies:
-      - vllm/
-      - tests/evals/gpt_oss/
-      - .buildkite/intel_jobs/test-intel.yaml
-    commands:
-      - >-
-        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'pip install "gpt-oss[eval]==0.0.5" &&
-        cd tests &&
-        pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
+        pytest -v -s quantization/test_auto_round.py'
  - label: "XPU compressed tensors FP8 test"
    depends_on:
      - image-build-xpu
@@ -193,4 +168,4 @@ steps:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'cd tests &&
-        pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
+        pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
@@ -28,6 +28,11 @@
        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
      }
    },
+    {
+        "dataset_name": "sharegpt",
+        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
+      }
+    },
    {
      "test_name": "serving_llama8B_tp1_random_128_128",
      "server_parameters": {
@@ -31,46 +31,8 @@ steps:
      - text: "What is the release version?"
        key: release-version

-  - group: "Build CUDA 13.0 Python wheels"
+  - group: "Build Python wheels"
    key: "build-wheels"
-    steps:
-      - label: "Build wheel - aarch64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-arm64-cuda-13-0
-        agents:
-          queue: arm64_cpu_queue_release
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - label: "Build wheel - x86_64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-x86-cuda-13-0
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
-        env:
-          DOCKER_BUILDKIT: "1"
-
-  - block: "Unblock to build additional Python wheels"
-    depends_on: ~
-    key: block-build-additional-wheels
-    if: build.env("NIGHTLY") != "1"
-
-  - group: "Build additional Python wheels"
-    key: "build-additional-wheels"
-    depends_on: block-build-additional-wheels
-    allow_dependency_failure: true
    steps:
      - label: "Build wheel - aarch64 - CUDA 12.9"
        depends_on: ~
@@ -86,6 +48,20 @@ steps:
        env:
          DOCKER_BUILDKIT: "1"

+      - label: "Build wheel - aarch64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-arm64-cuda-13-0
+        agents:
+          queue: arm64_cpu_queue_release
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+        env:
+          DOCKER_BUILDKIT: "1"
+
      - label: "Build wheel - aarch64 - CPU"
        depends_on: ~
        id: build-wheel-arm64-cpu
@@ -137,7 +113,7 @@ steps:
          - 'mv artifacts/reassembled/wheel "artifacts/dist/$$wheel_name"'
          - "aws sts get-caller-identity"
          - "VLLM_WHEEL_PLATFORM=macos bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        plugins:
          - aws-assume-role-with-web-identity#v1.6.0:
              role-arn: arn:aws:iam::936637512419:role/vllm-release-macos-wheel-uploader
@@ -157,6 +133,20 @@ steps:
        env:
          DOCKER_BUILDKIT: "1"

+      - label: "Build wheel - x86_64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-x86-cuda-13-0
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+        env:
+          DOCKER_BUILDKIT: "1"
+
      - label: "Build wheel - x86_64 - CPU"
        depends_on: ~
        id: build-wheel-x86-cpu
@@ -172,26 +162,12 @@ steps:
          DOCKER_BUILDKIT: "1"

  - label: "Generate and upload wheel indices"
-    key: generate-wheel-indices
    depends_on: "build-wheels"
    allow_dependency_failure: true
-    if: build.env("NIGHTLY") != "1"
    agents:
      queue: cpu_queue_release
    commands:
-      - "UPDATE_VERSION_INDEX=0 bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
-
-  - label: "Regenerate indices with additional wheels"
-    key: generate-additional-wheel-indices
-    depends_on:
-      - build-wheels
-      - build-additional-wheels
-      - generate-wheel-indices
-    allow_dependency_failure: true
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - 'UPDATE_NIGHTLY_INDEX="$${NIGHTLY:-0}" bash .buildkite/scripts/generate-and-upload-nightly-index.sh'
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

  - block: "Unblock to build release Docker images"
    depends_on: ~
@@ -590,370 +566,366 @@ steps:
  #
  # =============================================================================

-  - group: "Build ROCm Wheel / Image "
-    key: "build-rocm-wheel-image"
+  # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
+  - label: ":rocm: Build ROCm Base Image & Wheels"
+    id: build-rocm-base-wheels
    depends_on: ~
-    steps:
-      # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
-      - label: ":rocm: Build ROCm Base Image & Wheels"
-        id: build-rocm-base-wheels
-        depends_on: ~
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - |
-            set -euo pipefail
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - |
+        set -euo pipefail

-            # Generate cache key
-            CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-            ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
+        # Generate cache key
+        CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+        ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"

-            echo "========================================"
-            echo "ROCm Base Build Configuration"
-            echo "========================================"
-            echo "  CACHE_KEY: $${CACHE_KEY}"
-            echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
-            echo "========================================"
-            
-            # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            
-            IMAGE_EXISTS=false
-            WHEELS_EXIST=false
-            
-            # Check ECR for Docker image
+        echo "========================================"
+        echo "ROCm Base Build Configuration"
+        echo "========================================"
+        echo "  CACHE_KEY: $${CACHE_KEY}"
+        echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
+        echo "========================================"
+        
+        # Login to ECR
+        aws ecr-public get-login-password --region us-east-1 | \
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+        
+        IMAGE_EXISTS=false
+        WHEELS_EXIST=false
+        
+        # Check ECR for Docker image

-            if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
-              IMAGE_EXISTS=true
-              echo "ECR image cache HIT"
-            fi
-            
-            # Check S3 for wheels
-            WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
-            if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
-              WHEELS_EXIST=true
-              echo "S3 wheels cache HIT"
-            fi
+        if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
+          IMAGE_EXISTS=true
+          echo "ECR image cache HIT"
+        fi
+        
+        # Check S3 for wheels
+        WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
+        if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
+          WHEELS_EXIST=true
+          echo "S3 wheels cache HIT"
+        fi

-            
-            # Scenario 1: Both cached (best case)
-            if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
-              echo ""
-              echo "FULL CACHE HIT - Reusing both image and wheels"
-              echo ""
+        
+        # Scenario 1: Both cached (best case)
+        if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
+          echo ""
+          echo "FULL CACHE HIT - Reusing both image and wheels"
+          echo ""

-              # Download wheels
-              .buildkite/scripts/cache-rocm-base-wheels.sh download
-              
-              # Save ECR tag for downstream jobs
-              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-              
-            # Scenario 2: Full rebuild needed
-            else
-              echo ""
-              echo " CACHE MISS - Building from scratch..."
-              echo ""
-              
-              # Build full base image and push to ECR
-              DOCKER_BUILDKIT=1 docker buildx build \
-                --file docker/Dockerfile.rocm_base \
-                --tag "$${ECR_CACHE_TAG}" \
-                --build-arg USE_SCCACHE=1 \
-                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-                --build-arg SCCACHE_REGION_NAME=us-west-2 \
-                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-                --push \
-                .
-              
-              # Build wheel extraction stage
-              DOCKER_BUILDKIT=1 docker buildx build \
-                --file docker/Dockerfile.rocm_base \
-                --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
-                --target debs_wheel_release \
-                --build-arg USE_SCCACHE=1 \
-                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-                --build-arg SCCACHE_REGION_NAME=us-west-2 \
-                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-                --load \
-                .
-              
-              # Extract and upload wheels
-              mkdir -p artifacts/rocm-base-wheels
-              cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
-              docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
-              docker rm $${cid}
-              
-              .buildkite/scripts/cache-rocm-base-wheels.sh upload
+          # Download wheels
+          .buildkite/scripts/cache-rocm-base-wheels.sh download
+          
+          # Save ECR tag for downstream jobs
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+          
+        # Scenario 2: Full rebuild needed
+        else
+          echo ""
+          echo " CACHE MISS - Building from scratch..."
+          echo ""
+          
+          # Build full base image and push to ECR
+          DOCKER_BUILDKIT=1 docker buildx build \
+            --file docker/Dockerfile.rocm_base \
+            --tag "$${ECR_CACHE_TAG}" \
+            --build-arg USE_SCCACHE=1 \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --push \
+            .
+          
+          # Build wheel extraction stage
+          DOCKER_BUILDKIT=1 docker buildx build \
+            --file docker/Dockerfile.rocm_base \
+            --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
+            --target debs_wheel_release \
+            --build-arg USE_SCCACHE=1 \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --load \
+            .
+          
+          # Extract and upload wheels
+          mkdir -p artifacts/rocm-base-wheels
+          cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
+          docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
+          docker rm $${cid}
+          
+          .buildkite/scripts/cache-rocm-base-wheels.sh upload

-              # Cache base docker image to ECR
-              docker push "$${ECR_CACHE_TAG}"
-              
-              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-              
-              echo ""
-              echo " Build complete - Image and wheels cached"
-            fi
+          # Cache base docker image to ECR
+          docker push "$${ECR_CACHE_TAG}"
+          
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+          
+          echo ""
+          echo " Build complete - Image and wheels cached"
+        fi

-        artifact_paths:
-          - "artifacts/rocm-base-wheels/*.whl"
-        env:
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+    artifact_paths:
+      - "artifacts/rocm-base-wheels/*.whl"
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"

-      # ROCm Job 2: Build vLLM ROCm Wheel
-      - label: ":python: Build vLLM ROCm Wheel - x86_64"
-        id: build-rocm-vllm-wheel
-        depends_on:
-          - step: build-rocm-base-wheels
-            allow_failure: false
-        agents:
-          queue: cpu_queue_release
-        timeout_in_minutes: 180
-        commands:
-          # Download artifacts and prepare Docker image
-          - |
-            set -euo pipefail
+  # ROCm Job 2: Build vLLM ROCm Wheel
+  - label: ":python: Build vLLM ROCm Wheel - x86_64"
+    id: build-rocm-vllm-wheel
+    depends_on:
+      - step: build-rocm-base-wheels
+        allow_failure: false
+    agents:
+      queue: cpu_queue_release
+    timeout_in_minutes: 180
+    commands:
+      # Download artifacts and prepare Docker image
+      - |
+        set -euo pipefail

-            # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
-            # This fixes version detection when tags are moved/force-pushed
-            echo "Fetching latest tags from origin..."
-            git fetch --tags --force origin
-            
-            # Log tag information for debugging version detection
-            echo "========================================"
-            echo "Git Tag Verification"
-            echo "========================================"
-            echo "Current HEAD: $(git rev-parse HEAD)"
-            echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
-            echo ""
-            echo "Recent tags (pointing to commits near HEAD):"
-            git tag -l --sort=-creatordate | head -5
-            echo "setuptools_scm version detection:"
-            pip install -q setuptools_scm 2>/dev/null || true
-            python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
-            echo "========================================"
+        # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
+        # This fixes version detection when tags are moved/force-pushed
+        echo "Fetching latest tags from origin..."
+        git fetch --tags --force origin
+        
+        # Log tag information for debugging version detection
+        echo "========================================"
+        echo "Git Tag Verification"
+        echo "========================================"
+        echo "Current HEAD: $(git rev-parse HEAD)"
+        echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
+        echo ""
+        echo "Recent tags (pointing to commits near HEAD):"
+        git tag -l --sort=-creatordate | head -5
+        echo "setuptools_scm version detection:"
+        pip install -q setuptools_scm 2>/dev/null || true
+        python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
+        echo "========================================"

-            # Download wheel artifacts from current build
-            echo "Downloading wheel artifacts from current build"
-            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+        # Download wheel artifacts from current build
+        echo "Downloading wheel artifacts from current build"
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .

-            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-            if [ -z "$${ECR_IMAGE_TAG}" ]; then
-              echo "ERROR: rocm-base-image-tag metadata not found"
-              echo "This should have been set by the build-rocm-base-wheels job"
-              exit 1
-            fi
-            
-            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-            
-            # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            
-            # Pull base Docker image from ECR
-            docker pull "$${ECR_IMAGE_TAG}"
-            
-            echo "Loaded base image: $${ECR_IMAGE_TAG}"
-            
-            # Prepare base wheels for Docker build context
-            mkdir -p docker/context/base-wheels
-            touch docker/context/base-wheels/.keep
-            cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
-            echo "Base wheels for vLLM build:"
-            ls -lh docker/context/base-wheels/
+        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+        if [ -z "$${ECR_IMAGE_TAG}" ]; then
+          echo "ERROR: rocm-base-image-tag metadata not found"
+          echo "This should have been set by the build-rocm-base-wheels job"
+          exit 1
+        fi
+        
+        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+        
+        # Login to ECR
+        aws ecr-public get-login-password --region us-east-1 | \
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+        
+        # Pull base Docker image from ECR
+        docker pull "$${ECR_IMAGE_TAG}"
+        
+        echo "Loaded base image: $${ECR_IMAGE_TAG}"
+        
+        # Prepare base wheels for Docker build context
+        mkdir -p docker/context/base-wheels
+        touch docker/context/base-wheels/.keep
+        cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
+        echo "Base wheels for vLLM build:"
+        ls -lh docker/context/base-wheels/

-            echo "========================================"
-            echo "Building vLLM wheel with:"
-            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-            echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
-            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-            echo "========================================"
+        echo "========================================"
+        echo "Building vLLM wheel with:"
+        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+        echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
+        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+        echo "========================================"

-            # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
-            DOCKER_BUILDKIT=1 docker build \
-              --file docker/Dockerfile.rocm \
-              --target export_vllm_wheel_release \
-              --output type=local,dest=rocm-dist \
-              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-              --build-arg REMOTE_VLLM=0 \
-              --build-arg GIT_REPO_CHECK=1 \
-              --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 \
-              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-              .
-            echo "Built vLLM wheel:"
-            ls -lh rocm-dist/*.whl
-            # Copy wheel to artifacts directory
-            mkdir -p artifacts/rocm-vllm-wheel
-            cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
-            echo "Final vLLM wheel:"
-            ls -lh artifacts/rocm-vllm-wheel/
-        artifact_paths:
-          - "artifacts/rocm-vllm-wheel/*.whl"
-        env:
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+        # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
+        DOCKER_BUILDKIT=1 docker build \
+          --file docker/Dockerfile.rocm \
+          --target export_vllm_wheel_release \
+          --output type=local,dest=rocm-dist \
+          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+          --build-arg REMOTE_VLLM=0 \
+          --build-arg GIT_REPO_CHECK=1 \
+          --build-arg USE_SCCACHE=1 \
+          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+          --build-arg SCCACHE_REGION_NAME=us-west-2 \
+          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+          .
+        echo "Built vLLM wheel:"
+        ls -lh rocm-dist/*.whl
+        # Copy wheel to artifacts directory
+        mkdir -p artifacts/rocm-vllm-wheel
+        cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
+        echo "Final vLLM wheel:"
+        ls -lh artifacts/rocm-vllm-wheel/
+    artifact_paths:
+      - "artifacts/rocm-vllm-wheel/*.whl"
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"

-      # ROCm Job 3: Upload Wheels to S3
-      - label: ":s3: Upload ROCm Wheels to S3"
-        id: upload-rocm-wheels
-        depends_on:
-          - step: build-rocm-vllm-wheel
-            allow_failure: false
-        agents:
-          queue: cpu_queue_release
-        timeout_in_minutes: 60
-        commands:
-          # Download all wheel artifacts and run upload
-          - |
-            set -euo pipefail
+  # ROCm Job 3: Upload Wheels to S3
+  - label: ":s3: Upload ROCm Wheels to S3"
+    id: upload-rocm-wheels
+    depends_on:
+      - step: build-rocm-vllm-wheel
+        allow_failure: false
+    agents:
+      queue: cpu_queue_release
+    timeout_in_minutes: 60
+    commands:
+      # Download all wheel artifacts and run upload
+      - |
+        set -euo pipefail

-            # Download artifacts from current build
-            echo "Downloading artifacts from current build"
-            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-            buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+        # Download artifacts from current build
+        echo "Downloading artifacts from current build"
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+        buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .

-            # # Run upload script
-            bash .buildkite/scripts/upload-rocm-wheels.sh
-        env:
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+        # Run upload script
+        bash .buildkite/scripts/upload-rocm-wheels.sh
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"

-      # ROCm Job 4: Annotate ROCm Wheel Release
-      - label: ":memo: Annotate ROCm wheel release"
-        id: annotate-rocm-release
-        depends_on:
-          - upload-rocm-wheels
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/annotate-rocm-release.sh"
-        env:
-          S3_BUCKET: "vllm-wheels"
+  # ROCm Job 4: Annotate ROCm Wheel Release
+  - label: ":memo: Annotate ROCm wheel release"
+    id: annotate-rocm-release
+    depends_on:
+      - upload-rocm-wheels
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/annotate-rocm-release.sh"
+    env:
+      S3_BUCKET: "vllm-wheels"

-      # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
-      # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
-      # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
-      - block: "Generate Root Index for ROCm Wheels for Release"
-        key: block-generate-root-index-rocm-wheels
-        depends_on: upload-rocm-wheels
+  # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
+  # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
+  # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
+  - block: "Generate Root Index for ROCm Wheels for Release"
+    key: block-generate-root-index-rocm-wheels
+    depends_on: upload-rocm-wheels

-      - label: ":package: Generate Root Index for ROCm Wheels for Release"
-        depends_on: block-generate-root-index-rocm-wheels
-        id: generate-root-index-rocm-wheels
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
-        env:
-          S3_BUCKET: "vllm-wheels"
-          VARIANT: "rocm723"
+  - label: ":package: Generate Root Index for ROCm Wheels for Release"
+    depends_on: block-generate-root-index-rocm-wheels
+    id: generate-root-index-rocm-wheels
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
+    env:
+      S3_BUCKET: "vllm-wheels"
+      VARIANT: "rocm723"

-      # ROCm Job 6: Build ROCm Release Docker Image
-      - label: ":docker: Build release image - x86_64 - ROCm"
-        id: build-rocm-release-image
-        depends_on:
-          - step: block-build-release-images
-            allow_failure: true
-          - step: build-rocm-base-wheels
-            allow_failure: false
-        agents:
-          queue: cpu_queue_release
-        timeout_in_minutes: 60
-        commands:
-          - |
-            set -euo pipefail
-            
-            # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            
-            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-            if [ -z "$${ECR_IMAGE_TAG}" ]; then
-              echo "ERROR: rocm-base-image-tag metadata not found"
-              echo "This should have been set by the build-rocm-base-wheels job"
-              exit 1
-            fi
-            
-            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-            
-            # Pull base Docker image from ECR
-            docker pull "$${ECR_IMAGE_TAG}"
-            
-            echo "Loaded base image: $${ECR_IMAGE_TAG}"
-            
-            # Pass the base image ECR tag to downstream steps (nightly publish)
-            buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
-            
-            echo "========================================"
-            echo "Building vLLM ROCm release image with:"
-            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-            echo "========================================"
-            
-            # Build vLLM ROCm release image using cached base
-            DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
-              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-              --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 \
-              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
-              --target vllm-openai \
-              --progress plain \
-              -f docker/Dockerfile.rocm .
-            
-            # Push to ECR
-            docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
+  # ROCm Job 6: Build ROCm Release Docker Image
+  - label: ":docker: Build release image - x86_64 - ROCm"
+    id: build-rocm-release-image
+    depends_on:
+      - step: block-build-release-images
+        allow_failure: true
+      - step: build-rocm-base-wheels
+        allow_failure: false
+    agents:
+      queue: cpu_queue_release
+    timeout_in_minutes: 60
+    commands:
+      - |
+        set -euo pipefail
+        
+        # Login to ECR
+        aws ecr-public get-login-password --region us-east-1 | \
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+        
+        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+        if [ -z "$${ECR_IMAGE_TAG}" ]; then
+          echo "ERROR: rocm-base-image-tag metadata not found"
+          echo "This should have been set by the build-rocm-base-wheels job"
+          exit 1
+        fi
+        
+        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+        
+        # Pull base Docker image from ECR
+        docker pull "$${ECR_IMAGE_TAG}"
+        
+        echo "Loaded base image: $${ECR_IMAGE_TAG}"
+        
+        # Pass the base image ECR tag to downstream steps (nightly publish)
+        buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
+        
+        echo "========================================"
+        echo "Building vLLM ROCm release image with:"
+        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+        echo "========================================"
+        
+        # Build vLLM ROCm release image using cached base
+        DOCKER_BUILDKIT=1 docker build \
+          --build-arg max_jobs=16 \
+          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+          --build-arg USE_SCCACHE=1 \
+          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+          --build-arg SCCACHE_REGION_NAME=us-west-2 \
+          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+          --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
+          --target vllm-openai \
+          --progress plain \
+          -f docker/Dockerfile.rocm .
+        
+        # Push to ECR
+        docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm

-            echo ""
-            echo " Successfully built and pushed ROCm release image"
-            echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
-            echo ""
-        env:
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+        echo ""
+        echo " Successfully built and pushed ROCm release image"
+        echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
+        echo ""
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"

-      - label: "Publish nightly XPU image to DockerHub"
-        depends_on:
-          - create-manifest-xpu
-        if: build.env("NIGHTLY") == "1"
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
+  - label: "Publish nightly XPU image to DockerHub"
+    depends_on:
+      - create-manifest-xpu
+    if: build.env("NIGHTLY") == "1"
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
+    plugins:
+      - docker-login#v3.0.0:
+          username: vllmbot
+          password-env: DOCKERHUB_TOKEN
+    env:
+      DOCKER_BUILDKIT: "1"
+      DOCKERHUB_USERNAME: "vllmbot"

-      - label: "Publish nightly ROCm image to DockerHub"
-        depends_on:
-          - build-rocm-release-image
-        if: build.env("NIGHTLY") == "1"
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
-          # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
+  - label: "Publish nightly ROCm image to DockerHub"
+    depends_on:
+      - build-rocm-release-image
+    if: build.env("NIGHTLY") == "1"
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
+      # Clean up old nightly builds (keep only last 14)
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
+    plugins:
+      - docker-login#v3.0.0:
+          username: vllmbot
+          password-env: DOCKERHUB_TOKEN
+    env:
+      DOCKER_BUILDKIT: "1"
+      DOCKERHUB_USERNAME: "vllmbot"

  # =============================================================================
  # Publish to DockerHub and PyPI (at the end so all builds complete first)
@@ -1002,8 +974,6 @@ steps:
        depends_on:
          - input-release-version
          - build-wheels
-          - build-additional-wheels
-          - generate-additional-wheel-indices

      - label: "Upload release wheels to PyPI"
        depends_on:
@@ -7,9 +7,6 @@

 set -euo pipefail

-# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
-git fetch --tags --force origin
-
 # The Rust frontend build needs protoc.
 if ! command -v protoc >/dev/null 2>&1; then
  brew install protobuf
@@ -17,7 +17,7 @@ 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 tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
+DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
 DEFAULT_CI_BASE_METADATA_VERSION="1"
 IMAGE_EXISTED_BEFORE_BUILD=0

@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
    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.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
-    metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
+    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")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
    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.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
+    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:-}"

@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
        return 0
    fi

-    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
+    for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
        if [[ -n "${!var:-}" ]]; then
            echo "Using provided ${var}: ${!var}"
            continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
 compute_dependency_cache_keys() {
    local bake_dir=""
    local dockerfile_rocm=""
-    local nixl_branch=""
+    local rixl_branch=""
    local ucx_branch=""
    local rocshmem_branch=""
    local deepep_branch=""
-    local nixl_material=""
+    local rixl_material=""
    local rocshmem_material=""
    local deepep_material=""

    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
-    nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
+    rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
    ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
    rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
    deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")

-    if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
-        nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
-        NIXL_CACHE_KEY=$(
+    if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
+        rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
+        RIXL_CACHE_KEY=$(
            compose_dependency_cache_key \
-                "${nixl_branch}-ucx-${ucx_branch}" \
-                "${nixl_material}"
+                "${rixl_branch}-ucx-${ucx_branch}" \
+                "${rixl_material}"
        )
-        export NIXL_CACHE_KEY
-        echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
+        export RIXL_CACHE_KEY
+        echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
    fi

    if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"

    case "${target}" in
-        nixl-rocm-ci)
-            if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
-            elif [[ -n "${NIXL_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
+        rixl-rocm-ci)
+            if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
+            elif [[ -n "${RIXL_BRANCH:-}" ]]; then
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
            fi
            ;;
        rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {

 resolve_ci_base_dependency_targets() {
    local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
-    local nixl_ref=""
+    local rixl_ref=""
    local rocshmem_ref=""
    local deepep_ref=""

@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
    case "${mode}" in
        always)
            echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
-            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
+            for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
                if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
                    add_dependency_cache_target "${target}"
                fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
            ;;
    esac

-    if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
-        nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
-        if dependency_cache_ref_exists "${nixl_ref}"; then
-            echo "NIXL dependency cache exists: ${nixl_ref}"
+    if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
+        rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
+        if dependency_cache_ref_exists "${rixl_ref}"; then
+            echo "RIXL dependency cache exists: ${rixl_ref}"
        else
-            echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
-            add_dependency_cache_target "nixl-rocm-ci"
+            echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
+            add_dependency_cache_target "rixl-rocm-ci"
        fi
    fi

@@ -45,10 +45,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"

-# copy to /nightly/ only when enabled for a main branch build that is not a PR
-if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
-      "$BUILDKITE_BRANCH" == "main" && \
-      "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
+# copy to /nightly/ only if it is on the main branch and not a PR
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
    echo "Uploading indices to overwrite /nightly/"
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
 fi
@@ -69,7 +67,7 @@ pure_version="${version%%+*}"
 echo "Pure version (without variant): $pure_version"

 # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
+if [[ "$version" != *"dev"* ]]; then
    echo "Re-generating indices for /$pure_version/"
    rm -rf "${INDICES_OUTPUT_DIR:?}"
    mkdir -p "$INDICES_OUTPUT_DIR"
@@ -35,7 +35,7 @@ set -o pipefail
 : "${PY_COLORS:=1}"
 : "${ROCM_DOCKER_TTY:=1}"
 : "${PYTHONFAULTHANDLER:=1}"
-: "${PYTEST_TIMEOUT:=2400}"
+: "${PYTEST_TIMEOUT:=2100}"
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
 fi
@@ -45,9 +45,9 @@ fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
 fi
-# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
+# Dump stacks after 15 minutes, then stop an individual test after 35 minutes.
 if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
+  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
 fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
  " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
@@ -387,7 +387,6 @@ initialize_native_environment() {
  local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
  local job_id_suffix=""
  local native_root=""
-  local hf_fstype=""
  local hf_mount=""

  if [[ "$(id -u)" -ne 0 ]]; then
@@ -401,19 +400,16 @@ initialize_native_environment() {
  native_root="/tmp/vllm-native-${job_id}"
  TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
  VLLM_RPC_BASE_PATH="/tmp"
-  TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
-  TRITON_CACHE_DIR="${native_root}/cache/triton"
-  VLLM_CACHE_ROOT="${native_root}/cache/vllm"
-  XDG_CACHE_HOME="${native_root}/cache/xdg"
+  : "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
+  : "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
+  : "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
+  : "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
  : "${HF_HOME:=/home/buildkite-agent/huggingface}"
-  # datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
-  # Keep processed datasets job-local while retaining the persistent Hub cache.
-  HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
  : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
  : "${HF_HUB_ETAG_TIMEOUT:=60}"
  export TMPDIR VLLM_RPC_BASE_PATH
  export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
-  export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
+  export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
  export PYTORCH_ROCM_ARCH=""

  mkdir -p "${TMPDIR}" \
@@ -421,10 +417,7 @@ initialize_native_environment() {
    "${TRITON_CACHE_DIR}" \
    "${VLLM_CACHE_ROOT}" \
    "${XDG_CACHE_HOME}" \
-    "${HF_HOME}" \
-    "${HF_DATASETS_CACHE}" || return 1
-
-  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
+    "${HF_HOME}" || return 1

  if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
    if ! command -v findmnt >/dev/null 2>&1; then
@@ -437,18 +430,6 @@ initialize_native_environment() {
      return 1
    fi
  fi
-
-  if command -v findmnt >/dev/null 2>&1; then
-    hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true)
-  fi
-  if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then
-    # Keep hf-xet state local and avoid vectored writes on shared NFS.
-    export HF_XET_CACHE="${native_root}/cache/hf-xet"
-    export HF_XET_HIGH_PERFORMANCE=0
-    export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
-    mkdir -p "${HF_XET_CACHE}" || return 1
-    echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}"
-  fi
 }

 run_native_preflight() {
@@ -1,11 +1,10 @@
 #!/bin/bash
 set -euox pipefail

-export VLLM_CPU_KVCACHE_SPACE=1
+export VLLM_CPU_KVCACHE_SPACE=1 
 export VLLM_CPU_CI_ENV=1
-# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
-# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
-# when callers specify fullgraph=True.
+# Reduce sub-processes for acceleration
+export TORCH_COMPILE_DISABLE=1 
 export VLLM_ENABLE_V1_MULTIPROCESSING=0

 SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -50,15 +49,15 @@ wait_for_pid_and_check_log() {
 }

 # Test Sky Lake (AVX512F)
-./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
+./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
 PID_TEST_0=$!

 # Test Cascade Lake (AVX512F + VNNI)
-./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
+./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
 PID_TEST_1=$!

 # Test Cooper Lake (AVX512F + VNNI + BF16)
-./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
+./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
 PID_TEST_2=$!

 wait_for_pid_and_check_log $PID_TEST_0 test_0.log
@@ -40,9 +40,7 @@ function cpu_tests() {
    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
-    pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"

  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -99,4 +97,3 @@ function cpu_tests() {
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
 timeout 2h bash -c cpu_tests
-
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    -e CMDS \
    --name "${container_name}" \
    "${IMAGE}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
    >/dev/null
 } 9>/tmp/docker-pull.lock

@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"

-# Only scheduled nightly builds should update the moving nightly index.
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+# Update rocm/nightly/ if on main branch and not a PR
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Updating rocm/nightly/ index..."
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
 fi
@@ -147,7 +147,7 @@ echo ""
 echo "Install command (by commit):"
 echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
 echo ""
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Install command (nightly):"
    echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
 fi
@@ -16,9 +16,8 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 125
+      device: mi325_1
+      timeout_in_minutes: 95
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
  key: basic-correctness
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -18,8 +18,7 @@ steps:
  - pytest -v -s basic_correctness/test_cpu_offload.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi325_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
  key: benchmarks-cli-test
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -15,7 +15,6 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 40
      depends_on:
      - image-build-amd

@@ -16,10 +16,8 @@ steps:
  commands:
    - pytest -v -s cuda/test_cuda_context.py
    - pytest -v -s cuda/test_platform_no_cuda_init.py
-    - pytest -v -s cuda/test_cuda_compatibility_path.py

 - label: Cudagraph
-  device: h200_35gb
  key: cudagraph
  timeout_in_minutes: 30
  source_file_dependencies:
@@ -27,10 +25,7 @@ steps:
  - vllm/v1/cudagraph_dispatcher.py
  - vllm/config/compilation.py
  - vllm/compilation
-  - vllm/v1/worker/encoder_cudagraph.py
-  - vllm/v1/worker/encoder_cudagraph_defs.py
  commands:
    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
-    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
-    - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
@@ -17,7 +17,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 85
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -68,7 +68,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 40
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -94,7 +94,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 85
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -120,7 +120,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 55
+      timeout_in_minutes: 80
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -131,22 +131,6 @@ steps:
        - 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: NixlConnector PD edge case test (2 GPUs)
-  key: nixlconnector-pd-edge-cases-2-gpus
-  timeout_in_minutes: 40
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-    - vllm/v1/core/sched/
-    - tests/v1/kv_connector/nixl_integration/
-  env:
-    PREFILL_GPU_ID: "0"
-    DECODE_GPU_ID: "1"
-  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
-
 - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
  key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
  timeout_in_minutes: 25
@@ -193,7 +177,7 @@ steps:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 45
+      timeout_in_minutes: 70
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -41,7 +41,6 @@ steps:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -28,9 +28,8 @@ steps:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 40
+      device: mi325_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd

@@ -40,21 +39,19 @@ steps:
  source_file_dependencies:
    - vllm/v1/engine/
    - tests/v1/engine/
-    - tests/v1/test_tensor_ipc_queue.py
  commands:
    - pytest -v -s v1/engine/test_preprocess_error_handling.py
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
-    - pytest -v -s v1/test_tensor_ipc_queue.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 45
+      device: mi325_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd

 - label: e2e Scheduling (1 GPU)
  key: e2e-scheduling-1-gpu
-  timeout_in_minutes: 53
+  timeout_in_minutes: 35
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/
@@ -63,8 +60,8 @@ steps:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd

@@ -79,8 +76,8 @@ steps:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 50
+      device: mi325_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -119,7 +116,6 @@ steps:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd

@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Entrypoints Unit Tests
-  device: h200_35gb
  key: entrypoints-unit-tests
  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
@@ -16,7 +15,6 @@ steps:
  - pytest -v -s entrypoints/weight_transfer

 - label: Entrypoints Integration (LLM)
-  device: h200_35gb
  key: entrypoints-integration-llm
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
@@ -30,16 +28,16 @@ steps:
  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server)
  key: entrypoints-integration-api-server
  device: h200_35gb
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -52,16 +50,13 @@ steps:
  - pytest -v -s entrypoints/scale_out
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server OpenAI - Part 1)
-  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-1
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -72,16 +67,14 @@ steps:
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server OpenAI - Part 2)
-  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-2
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -93,14 +86,12 @@ steps:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 70
+      device: mi325_1
+      timeout_in_minutes: 80
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server Generate)
-  device: h200_35gb
  key: entrypoints-integration-api-server-generate
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -117,14 +108,12 @@ steps:
  - pytest -v -s entrypoints/anthropic
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (Responses API)
-  device: h200_35gb
  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -159,9 +148,8 @@ steps:
  - pytest -v -s entrypoints/multimodal

 - label: Entrypoints Integration (Pooling)
-  device: h200_35gb
  key: entrypoints-integration-pooling
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -181,9 +169,7 @@ steps:
  - pytest -s entrypoints/openai/correctness/
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 30
+      device: mi325_1
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -18,7 +18,6 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -52,5 +51,4 @@ steps:
  - vllm/compilation/
  - tests/distributed/
  commands:
-  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
  - pytest -v -s distributed/test_elastic_ep.py
@@ -1,26 +0,0 @@
-group: Fault Tolerance
-depends_on:
-  - image-build
-steps:
- label: Fault Tolerance E2E (2xH100)
-  key: fault-tolerance-e2e-2xh100
-  timeout_in_minutes: 35
-  device: h100
-  num_devices: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/fault_tolerance/
-  - vllm/v1/worker/sentinel/
-  - vllm/entrypoints/serve/fault_tolerance/
-  - vllm/distributed/elastic_ep/
-  - vllm/distributed/device_communicators/
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/v1/fault_tolerance/
-  - tests/v1/distributed/test_external_lb_dp.py
-  commands:
-  # Base image has no nixl; install it or has_nixl_ep() skips the tests.
-  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
@@ -15,7 +15,6 @@ steps:
    - pytest -v -s tests/kernels/ir

 - label: Kernels Core Operation Test
-  device: h200_35gb
  key: kernels-core-operation-test
  timeout_in_minutes: 120
  source_file_dependencies:
@@ -61,45 +60,9 @@ steps:
  source_file_dependencies:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
  - vllm/models/deepseek_v4/common/ops/
-  - vllm/models/deepseek_v4/nvidia/
  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-  - tests/models/test_deepseek_v4_mega_moe.py
  commands:
    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
-    - pytest -v -s models/test_deepseek_v4_mega_moe.py
-
-# Catch-all for test files at the tests/kernels root. This job collects
-# the whole root so new files are wired by default.
-# Files with dedicated jobs elsewhere in this file are excluded via --ignore
-# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in
-# their own jobs / Kernels (B200)).
- label: Kernels Root Misc Test (B200)
-  key: kernels-root-misc-test-b200
-  timeout_in_minutes: 45
-  device: b200-k8s
-  source_file_dependencies:
-  - csrc/
-  - vllm/
-  - tests/kernels/
-  commands:
-    - pytest -v -s kernels/
-        --ignore=kernels/attention
-        --ignore=kernels/core
-        --ignore=kernels/helion
-        --ignore=kernels/ir
-        --ignore=kernels/mamba
-        --ignore=kernels/moe
-        --ignore=kernels/quantization
-        --ignore=kernels/test_concat_mla_q.py
-        --ignore=kernels/test_fused_qk_norm_rope_gate.py
-        --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-        --ignore=kernels/test_top_k_per_row.py
-        --ignore=kernels/test_kda.py
-        --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
-        --ignore=kernels/test_ll_bf16_gemm.py
-        --ignore=kernels/test_shuffle_rows.py
-    # BROKEN on main, pending kernel fixes (B200):
-    #   test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)

 - label: Kernels Attention Test %N
  key: kernels-attention-test
@@ -116,8 +79,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
@@ -155,9 +117,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 120
+      device: mi325_1
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
@@ -187,9 +147,8 @@ steps:
  parallelism: 5
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 65
      source_file_dependencies:
      - csrc/quantization/cutlass_w8a8/moe/
      - csrc/moe/
@@ -204,7 +163,6 @@ steps:
      - image-build-amd

 - label: Kernels Mamba Test
-  device: h200_35gb
  key: kernels-mamba-test
  timeout_in_minutes: 40
  source_file_dependencies:
@@ -214,6 +172,17 @@ steps:
  commands:
    - pytest -v -s kernels/mamba

+- label: Kernels KDA Test
+  timeout_in_minutes: 25
+  device: h200_18gb
+  source_file_dependencies:
+  - vllm/third_party/flash_linear_attention/ops/kda.py
+  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
+  - vllm/third_party/flash_linear_attention/ops/l2norm.py
+  - tests/kernels/test_kda.py
+  commands:
+    - pytest -v -s kernels/test_kda.py
+
 - label: Kernels DeepGEMM Test (H100)
  key: kernels-deepgemm-test-h100
  timeout_in_minutes: 35
@@ -266,11 +235,6 @@ steps:
  - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
-  - vllm/cute_utils/
-  - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
-  - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
-  - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
-  - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
  - tests/kernels/test_ll_bf16_gemm.py
  - tests/kernels/test_top_k_per_row.py
  commands:
@@ -300,8 +264,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
    - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
-    - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
-    - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
    - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
    # e2e
    - pytest -v -s tests/models/quantization/test_nvfp4.py
@@ -14,9 +14,8 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 45
+      device: mi325_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -79,28 +78,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt

- label: LM Eval PCP (4xB200)
-  key: lm-eval-pcp-4xb200
-  timeout_in_minutes: 360
-  device: b200-k8s
-  num_devices: 4
-  optional: true
-  source_file_dependencies:
-  - csrc/
-  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
-  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
-  - tests/evals/gsm8k/configs/models-pcp.txt
-  - vllm/model_executor/layers/quantization
-  - vllm/config/parallel.py
-  - vllm/distributed/parallel_state.py
-  - vllm/model_executor/layers/attention/mla_attention.py
-  - vllm/model_executor/layers/attention/pcp.py
-  - vllm/v1/worker/gpu/model_runner.py
-  - vllm/v1/worker/gpu/pcp_manager.py
-  autorun_on_main: true
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
-
 - label: LM Eval Large Models EP (2xB200)
  key: lm-eval-large-models-ep-2xb200
  timeout_in_minutes: 60
@@ -142,7 +119,7 @@ steps:
    amd:
      dind: false
      device: mi300_8
-      timeout_in_minutes: 40
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      commands:
@@ -337,7 +314,7 @@ steps:

 - label: LM Eval KV-Offload (2xH100)
  key: kv-offload-medium
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: h100
  num_devices: 2
  source_file_dependencies:
@@ -347,7 +324,7 @@ steps:
  - vllm/v1/simple_kv_offload/
  - tests/evals/gsm8k/test_gsm8k_offloading.py
  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
+  - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"

 - label: LM Eval KV-Offload (4xH100)
  key: kv-offload-large
@@ -14,10 +14,9 @@ steps:
  parallelism: 4
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 85
+      timeout_in_minutes: 65
      source_file_dependencies:
      - vllm/lora
      - tests/lora
@@ -47,4 +46,4 @@ steps:
    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
    - pytest -v -s -x lora/test_gptoss_tp.py
-    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -25,13 +25,13 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 50
+      timeout_in_minutes: 75
      depends_on:
      - image-build-amd

 - label: V1 Sample + Logits
  key: v1-sample-logits
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
    - vllm/config/
@@ -59,16 +59,13 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 70
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: V1 Core + KV + Metrics
-  device: h200_35gb
  key: v1-core-kv-metrics
-  timeout_in_minutes: 80
+  timeout_in_minutes: 60
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -92,7 +89,6 @@ steps:
    - tests/v1/kv_offload
    - tests/v1/simple_kv_offload
    - tests/v1/worker
-    - tests/v1/streaming_input
    - tests/v1/kv_connector/unit
    - tests/v1/ec_connector/unit
    - tests/v1/metrics
@@ -106,7 +102,6 @@ steps:
    - pytest -v -s v1/kv_offload
    - pytest -v -s v1/simple_kv_offload
    - pytest -v -s v1/worker
-    - pytest -v -s v1/streaming_input
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
@@ -115,9 +110,8 @@ steps:
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
+      timeout_in_minutes: 75
      depends_on:
      - image-build-amd

@@ -148,8 +142,6 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/core
    - pytest -v -s v1/structured_output
    - pytest -v -s v1/test_serial_utils.py
-    - pytest -v -s v1/test_kv_cache_spec_registry.py
-    - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'cpu_test' v1/metrics

@@ -213,7 +205,7 @@ steps:
  - vllm/multimodal
  - examples/
  commands:
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
     # for basic
    - python3 basic/offline_inference/chat.py
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -237,9 +229,7 @@ steps:
    - 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:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 75
+      device: mi325_1
      source_file_dependencies:
      - vllm/entrypoints
      - vllm/multimodal
@@ -266,7 +256,6 @@ steps:
  - vllm/utils/
  - vllm/v1/
  - tests/v1/tracing
-  - tests/tracing/
  commands:
  - "pip install \
      'opentelemetry-sdk>=1.26.0' \
@@ -274,14 +263,12 @@ steps:
      'opentelemetry-exporter-otlp>=1.26.0' \
      'opentelemetry-semantic-conventions-ai>=0.4.1'"
  - pytest -v -s v1/tracing
-  - pytest -v -s tracing
  mirror:
    amd:
-      dind: false
-      device: mi300_2
-      timeout_in_minutes: 30
+      device: mi325_2
      depends_on:
      - image-build-amd
+      optional: true

 - label: Python-only Installation
  key: python-only-installation
@@ -296,9 +283,8 @@ steps:
  - bash standalone_tests/python_only_compile.sh
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
-      soft_fail: true
+      device: mi325_1
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -398,7 +384,7 @@ steps:

 - label: Batch Invariance (A100)
  key: batch-invariance-a100
-  timeout_in_minutes: 60
+  timeout_in_minutes: 40
  device: a100
  source_file_dependencies:
    - vllm/v1/attention
@@ -408,11 +394,11 @@ steps:
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]

 - label: Batch Invariance (H100)
  key: batch-invariance-h100
-  timeout_in_minutes: 60
+  timeout_in_minutes: 40
  device: h100
  source_file_dependencies:
    - vllm/v1/attention
@@ -423,12 +409,12 @@ steps:
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]

 - label: Batch Invariance (B200)
  key: batch-invariance-b200
-  timeout_in_minutes: 45
+  timeout_in_minutes: 35
  device: b200-k8s
  source_file_dependencies:
    - vllm/v1/attention
@@ -439,14 +425,11 @@ steps:
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
-    - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
-    - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-    - pytest -v -s v1/determinism/test_online_batch_invariance.py
-
+  
 - label: Acceptance Length Test (Large Models) # optional
  device: h200_35gb
  key: acceptance-length-test-large-models
@@ -3,16 +3,13 @@ depends_on:
  - image-build
 steps:
 - label: Model Executor
-  device: h200_35gb
  key: model-executor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
  - vllm/model_executor
-  - vllm/model_executor/warmup
  - tests/model_executor
-  - tests/model_executor/test_jit_warmup.py
  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
@@ -30,16 +27,13 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
      - vllm/engine/arg_utils.py
      - vllm/config/model.py
      - vllm/model_executor
-      - vllm/model_executor/warmup
      - tests/model_executor
-      - tests/model_executor/test_jit_warmup.py
      - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
@@ -41,7 +41,7 @@ steps:
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
    - python3 basic/offline_inference/chat.py # for basic
    - 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  # TODO
@@ -42,39 +42,10 @@ steps:
    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 50
+      device: mi325_1
      depends_on:
      - image-build-amd

- label: Inkling Unit Tests (B200)
-  key: inkling-unit-tests-b200
-  timeout_in_minutes: 40
-  device: b200-k8s
-  source_file_dependencies:
-  - vllm/models/inkling/
-  - vllm/cute_utils/
-  - cmake/external_projects/tml_fa4.cmake
-  - tests/models/inkling/
-  commands:
-    # FA4 kernel tests require SM100; the suite skips them elsewhere.
-    - pytest -v -s models/inkling
-
- label: Kimi K3 Unit Tests (B200)
-  key: kimi-k3-unit-tests-b200
-  timeout_in_minutes: 40
-  device: b200-k8s
-  source_file_dependencies:
-  - vllm/models/kimi_k3/
-  - csrc/libtorch_stable/kimi_k3/
-  - tests/models/kimi_k3/
-  - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
-  - tests/kernels/test_bf16_skinny_gemm.py
-  commands:
-    # The native NVIDIA Kimi K3 kernels require the SM100 family.
-    - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
-
 - label: Basic Models Test (Other CPU) # 5min
  key: basic-models-test-other-cpu
  depends_on:
@@ -84,8 +55,7 @@ steps:
  - vllm/
  - tests/models/test_utils.py
  - tests/models/test_vision.py
-  - tests/models/test_adapters.py
  - tests/models/transformers/fusers/
  device: cpu-small
  commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
@@ -17,12 +17,10 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 45
      depends_on:
      - image-build-amd

 - label: Language Models Tests (Extra Standard) %N
-  device: h200_35gb
  key: language-models-tests-extra-standard
  timeout_in_minutes: 40
  source_file_dependencies:
@@ -40,7 +38,6 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -54,8 +51,8 @@ steps:
      - tests/models/language/pooling/test_classification.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
+
 - label: Language Models Tests (Hybrid) %N
-  device: h200_35gb
  key: language-models-tests-hybrid
  timeout_in_minutes: 65
  source_file_dependencies:
@@ -63,16 +60,16 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
+    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    # Shard the hybrid language model tests that are numerically stable on Hopper.
-    - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+    # Shard hybrid language model tests
+    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi325_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
      commands:
@@ -80,20 +77,6 @@ steps:
      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
      - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

-# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
-# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
-# correctness test on L4 until its H200 output matches the Transformers reference.
- label: Language Models Tests (Granite L4 Compatibility)
-  key: language-models-tests-granite-l4-compatibility
-  timeout_in_minutes: 65
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
-    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
-
 - label: Language Models Test (Extended Generation) # 80min
  device: h200_35gb
  key: language-models-test-extended-generation
@@ -104,6 +87,7 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
+    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
@@ -131,15 +115,14 @@ steps:
    - pytest -v -s models/language/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 95
+      device: mi325_1
+      timeout_in_minutes: 120
      depends_on:
      - image-build-amd

 - label: Language Models Test (MTEB)
  key: language-models-test-mteb
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  optional: true
  source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -14,15 +14,13 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -33,9 +31,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -51,15 +47,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 4: other + whisper"
  device: h200_35gb
  key: multi-modal-models-standard-4-other-whisper
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -70,9 +65,7 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 50
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -92,7 +85,7 @@ steps:

 - label: Multi-Modal Processor # 44min
  key: multi-modal-processor
-  timeout_in_minutes: 98
+  timeout_in_minutes: 65
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -116,7 +109,6 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -127,7 +119,6 @@ steps:
      - vllm/model_executor/model_loader/

 - label: Multi-Modal Models (Extended Generation 1)
-  device: h200_35gb
  key: multi-modal-models-extended-generation-1
  optional: true
  source_file_dependencies:
@@ -139,9 +130,7 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 90
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -176,9 +165,8 @@ steps:
    - pytest -v -s models/multimodal/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi325_1
+      timeout_in_minutes: 75
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -5,7 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
  device: h200_35gb
  key: pytorch-compilation-unit-tests
-  timeout_in_minutes: 150
+  timeout_in_minutes: 90
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -107,11 +107,17 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
+      depends_on:
+      - image-build-amd

 - label: PyTorch Fullgraph Smoke Test
-  device: h200_35gb
  key: pytorch-fullgraph-smoke-test
-  timeout_in_minutes: 90
+  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/__init__.py
  - vllm/_aiter_ops.py
@@ -143,42 +149,7 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
-
-# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
-# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
-# those SM90 cases skip while the architecture-compatible cases still run.
- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
-  key: pytorch-fullgraph-cudagraph-l4-compatibility
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/__init__.py
-  - vllm/_aiter_ops.py
-  - vllm/_custom_ops.py
-  - vllm/compilation/
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/env_override.py
-  - vllm/envs.py
-  - vllm/forward_context.py
-  - vllm/inputs/
-  - vllm/ir/
-  - vllm/kernels/
-  - vllm/logger.py
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/plugins/
-  - vllm/sampling_params.py
-  - vllm/sequence.py
-  - vllm/transformers_utils/
-  - vllm/triton_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - tests/compile
-  commands:
-  - pytest -s -v compile/fullgraph/test_full_cudagraph.py
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"

 - label: PyTorch Fullgraph
  key: pytorch-fullgraph
@@ -229,7 +200,6 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -3,11 +3,8 @@ depends_on:
  - image-build
 steps:
 - label: Quantization
-  device: h200_35gb
  key: quantization
-  timeout_in_minutes: 75
-  env:
-    VLLM_USE_V2_MODEL_RUNNER: "0"
+  timeout_in_minutes: 60
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -22,12 +19,9 @@ steps:
  # TODO(jerryzh168): resolve the above comment
  - uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
  - uv pip install --system conch-triton-kernels
-  # The SM90-only checkpoint currently contains a removed weight_chan_scale
-  # parameter. It was not exercised by the previous L4 job.
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py

 - label: Quantized Fusions
-  device: h200_35gb
  key: quantized-fusions
  timeout_in_minutes: 20
  source_file_dependencies:
@@ -58,11 +52,8 @@ steps:
    - pytest -s -v tests/quantization/test_blackwell_moe.py

 - label: Quantized Models Test
-  device: h200_35gb
  key: quantized-models-test
-  timeout_in_minutes: 65
-  env:
-    VLLM_USE_V2_MODEL_RUNNER: "0"
+  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
@@ -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 and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
  - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"

  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
@@ -81,7 +81,6 @@ steps:
  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

 - label: Rust Frontend Tool Use
-  device: h200_35gb
  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -19,18 +19,8 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 40
+      device: mi325_1
      depends_on:
      - image-build-amd
-      source_file_dependencies:
-      - vllm/model_executor/layers
-      - vllm/sampling_metadata.py
-      - vllm/v1/sample/
-      - vllm/entrypoints/generate/beam_search/
-      - tests/samplers
-      - tests/conftest.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
      commands:
      - pytest -v -s samplers
@@ -14,9 +14,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -54,9 +53,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 75
+      device: mi325_1
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -90,15 +88,14 @@ steps:
    - vllm/v1/spec_decode/
    - vllm/v1/worker/gpu/spec_decode/
    - tests/v1/e2e/spec_decode/
-    - tests/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
-    - python3 spec_decode/test_custom_proposer.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 35
+      device: mi325_1
+      timeout_in_minutes: 55
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -122,8 +119,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
@@ -174,31 +170,3 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
-
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
-  key: spec-decode-deepseek-mtp-parallel-load-b200
-  timeout_in_minutes: 30
-  device: b200-k8s
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/v1/spec_decode/llm_base_proposer.py
-    - vllm/v1/spec_decode/eagle.py
-    - vllm/v1/worker/gpu/spec_decode/eagle/
-    - vllm/model_executor/models/deepseek_mtp.py
-    - vllm/model_executor/models/deepseek_v2.py
-    - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
-  commands:
-    - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
-
- label: Spec Decode Acceptance Rates Nightly
-  key: spec-decode-acceptance-rates-nightly
-  timeout_in_minutes: 60
-  device: h200_35gb
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/v1/worker/gpu/spec_decode/
-    - tests/v1/e2e/spec_decode/
-  commands:
-    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
@@ -1,14 +0,0 @@
-group: Torch ABI
-depends_on:
-  - image-build
-steps:
- label: Torch Stable ABI Audit
-  key: torch-stable-abi-audit
-  timeout_in_minutes: 5
-  source_file_dependencies:
-  - .buildkite/check-torch-abi.py
-  - csrc/
-  - cmake/
-  - setup.py
-  commands:
-  - python3 /vllm-workspace/.buildkite/check-torch-abi.py
@@ -17,7 +17,6 @@ steps:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      commands:
@@ -47,7 +47,6 @@

 # Rust Frontend
 /rust/ @BugenZhao @njhill
-/rust/src/bench @esmeetu
 /build_rust.sh @BugenZhao @njhill
 /rust-toolchain.toml @BugenZhao @njhill
 /.buildkite/test_areas/rust* @BugenZhao @njhill
@@ -19,7 +19,6 @@ pull_request_rules:
  description: Comment on PR when pre-commit check fails
  conditions:
    - check-failure=pre-commit
-    - -check-cancelled=pre-commit
    - -closed
    - -draft
    - or:
@@ -182,18 +181,6 @@ pull_request_rules:
      add:
        - performance

- name: label-quantization
-  description: Automatically apply quantization label
-  conditions:
-    - label != stale
-    - or:
-      - files~=^vllm/model_executor/layers/quantization/
-      - title~=(?i)quant
-  actions:
-    label:
-      add:
-        - quantization
-
 - name: label-qwen
  description: Automatically apply qwen label
  conditions:
@@ -233,31 +220,6 @@ pull_request_rules:
      add:
        - gpt-oss

- name: label-kimi
-  description: Automatically apply kimi label
-  conditions:
-    - label != stale
-    - or:
-      - files~=(?i)kimi
-      - files~=(?i)moonshot
-      - title~=(?i)(?:kimi|moonshot)
-  actions:
-    label:
-      add:
-        - kimi
-
- name: label-k3
-  description: Automatically apply k3 label (launch triage; retire after ramp-down)
-  conditions:
-    - label != stale
-    - or:
-      - files~=(?i)kimi[-_]?k3
-      - title~=(?i)(?:kimi[-\s]?k3|\bk3\b)
-  actions:
-    label:
-      add:
-        - k3
-
 - name: label-nvidia
  description: Automatically apply nvidia label
  conditions:
@@ -130,66 +130,6 @@ jobs:
                  },
                ],
              },
-              kimi: {
-                keywords: [
-                  { term: "Kimi",     searchIn: "both" },
-                  { term: "Moonshot", searchIn: "both" },
-                ],
-                substrings: [
-                  { term: "moonshotai/", searchIn: "both" },
-                  { term: "kimi",        searchIn: "title" },
-                ],
-              },
-              k3: {
-                keywords: [
-                  { term: "Kimi K3", searchIn: "both" },
-                  { term: "K3",      searchIn: "title" },
-                ],
-                substrings: [
-                  { term: "moonshotai/kimi-k3", searchIn: "both" },
-                ],
-              },
-              quantization: {
-                keywords: [
-                  {
-                    term: "quantization",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "quantized",
-                    searchIn: "both"
-                  },
-                ],
-              },
-              "intel-gpu": {
-                // Keyword search - matches whole words only (with word boundaries)
-                keywords: [
-                  {
-                    term: "B50",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "B60",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "B70",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "intel gpu",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "Arc GPU",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "BMG",
-                    searchIn: "both"
-                  },
-                ],
-              },
              // Add more label configurations here as needed
              // example: {
              //   keywords: [...],
@@ -551,4 +491,4 @@ jobs:
              issue_number: context.issue.number,
              body: message,
            });
-            core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+            core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
@@ -80,9 +80,9 @@ jobs:
                  '',
                  '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.',
                  '',
-                  'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.',
+                  'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
                  '',
-                  'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.',
+                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
                  '',
                  'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
                  '',
@@ -41,7 +41,7 @@ jobs:
          if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) {
            core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`);
          } else {
-            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`);
+            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`);
          }

  pre-commit:
@@ -1,40 +0,0 @@
-name: Run CI from PR comment
-
-on:
-  issue_comment:
-    types: [created]
-
-concurrency:
-  group: run-ci-comment-${{ github.event.issue.number }}
-  cancel-in-progress: false
-
-permissions:
-  contents: read
-  issues: write
-  pull-requests: write
-
-jobs:
-  run-ci-command:
-    if: >-
-      github.event.issue.pull_request &&
-      (github.event.comment.body == '/ci run' ||
-      github.event.comment.body == '/ci retry')
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-        with:
-          ref: ${{ github.event.repository.default_branch }}
-          persist-credentials: false
-      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
-        with:
-          python-version: "3.12"
-      - name: Authorize and run CI command
-        run: >-
-          uv run --no-project --python 3.12
-          .github/workflows/scripts/run_ci_command.py
-        env:
-          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
-          BUILDKITE_ORGANIZATION: vllm
-          BUILDKITE_PIPELINE: ci
-          CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }}
-          GH_TOKEN: ${{ github.token }}
@@ -1,581 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-import os
-import sys
-import urllib.error
-import urllib.parse
-import urllib.request
-from collections.abc import Mapping, Sequence
-from typing import Any
-
-COMMAND_RUN_CI = "/ci run"
-COMMAND_RETRY_FAILED = "/ci retry"
-READY_LABELS = {"ready", "ready-run-all-tests"}
-TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
-ACTIVE_BUILD_STATES = {
-    "blocked",
-    "creating",
-    "scheduled",
-    "running",
-    "failing",
-    "canceling",
-    "waiting",
-    "waiting_failed",
-}
-RETRY_STATES = "failed,timed_out,expired"
-
-
-class ApiError(RuntimeError):
-    def __init__(self, status: int | None, message: str) -> None:
-        super().__init__(message)
-        self.status = status
-
-
-class HttpTransport:
-    def request(
-        self,
-        url: str,
-        *,
-        body: Mapping[str, Any] | None = None,
-        headers: Mapping[str, str] | None = None,
-        method: str = "GET",
-    ) -> Any:
-        data = None if body is None else json.dumps(body).encode()
-        request = urllib.request.Request(
-            url,
-            data=data,
-            headers=dict(headers or {}),
-            method=method,
-        )
-        try:
-            with urllib.request.urlopen(request, timeout=30) as response:
-                response_body = response.read().decode()
-        except urllib.error.HTTPError as error:
-            response_body = error.read().decode()
-            message = self._error_message(response_body, error.reason)
-            raise ApiError(
-                error.code,
-                f"API returned {error.code}: {message}",
-            ) from error
-        except urllib.error.URLError as error:
-            raise ApiError(None, f"API request failed: {error.reason}") from error
-
-        if not response_body:
-            return None
-        try:
-            return json.loads(response_body)
-        except json.JSONDecodeError as error:
-            raise ApiError(None, "API returned a non-JSON response.") from error
-
-    @staticmethod
-    def _error_message(response_body: str, fallback: str) -> str:
-        try:
-            parsed = json.loads(response_body)
-        except json.JSONDecodeError:
-            return fallback
-        return str(parsed.get("message", fallback))
-
-
-class GitHubClient:
-    def __init__(
-        self,
-        token: str,
-        repository: str,
-        transport: HttpTransport | None = None,
-    ) -> None:
-        if not token:
-            raise RuntimeError("GH_TOKEN is not set.")
-        self.owner, self.repo = repository.split("/", maxsplit=1)
-        self.transport = transport or HttpTransport()
-        self.headers = {
-            "Accept": "application/vnd.github+json",
-            "Authorization": f"Bearer {token}",
-            "Content-Type": "application/json",
-            "User-Agent": "vllm-ci-command",
-            "X-GitHub-Api-Version": "2022-11-28",
-        }
-
-    def _request(
-        self,
-        path: str,
-        *,
-        body: Mapping[str, Any] | None = None,
-        method: str = "GET",
-    ) -> Any:
-        return self.transport.request(
-            f"https://api.github.com{path}",
-            body=body,
-            headers=self.headers,
-            method=method,
-        )
-
-    def _repo_path(self, suffix: str) -> str:
-        owner = urllib.parse.quote(self.owner, safe="")
-        repo = urllib.parse.quote(self.repo, safe="")
-        return f"/repos/{owner}/{repo}{suffix}"
-
-    def _paginate(self, path: str) -> list[dict[str, Any]]:
-        results: list[dict[str, Any]] = []
-        separator = "&" if "?" in path else "?"
-        for page in range(1, 101):
-            response = self._request(f"{path}{separator}per_page=100&page={page}")
-            if not isinstance(response, list):
-                raise ApiError(None, "GitHub API returned an invalid list response.")
-            results.extend(response)
-            if len(response) < 100:
-                return results
-        raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")
-
-    def get_pr(self, number: int) -> dict[str, Any]:
-        return self._request(self._repo_path(f"/pulls/{number}"))
-
-    def get_permission(self, actor: str) -> str:
-        username = urllib.parse.quote(actor, safe="")
-        try:
-            response = self._request(
-                self._repo_path(f"/collaborators/{username}/permission")
-            )
-        except ApiError as error:
-            if error.status == 404:
-                return "none"
-            raise
-        return str(response["permission"])
-
-    def get_review_decision(self, number: int) -> str | None:
-        query = """
-          query($owner: String!, $repo: String!, $number: Int!) {
-            repository(owner: $owner, name: $repo) {
-              pullRequest(number: $number) {
-                reviewDecision
-              }
-            }
-          }
-        """
-        response = self._request(
-            "/graphql",
-            body={
-                "query": query,
-                "variables": {
-                    "number": number,
-                    "owner": self.owner,
-                    "repo": self.repo,
-                },
-            },
-            method="POST",
-        )
-        return response["data"]["repository"]["pullRequest"]["reviewDecision"]
-
-    def list_reviews(self, number: int) -> list[dict[str, Any]]:
-        return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))
-
-    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
-        return self._paginate(
-            self._repo_path(f"/issues/comments/{comment_id}/reactions")
-        )
-
-    def add_reaction(self, comment_id: int, content: str) -> None:
-        self._request(
-            self._repo_path(f"/issues/comments/{comment_id}/reactions"),
-            body={"content": content},
-            method="POST",
-        )
-
-    def add_comment(self, issue_number: int, body: str) -> None:
-        self._request(
-            self._repo_path(f"/issues/{issue_number}/comments"),
-            body={"body": body},
-            method="POST",
-        )
-
-
-class BuildkiteClient:
-    def __init__(
-        self,
-        token: str,
-        organization: str,
-        pipeline: str,
-        transport: HttpTransport | None = None,
-    ) -> None:
-        self.token = token
-        self.transport = transport or HttpTransport()
-        organization = urllib.parse.quote(organization, safe="")
-        pipeline = urllib.parse.quote(pipeline, safe="")
-        self.base_url = (
-            "https://api.buildkite.com/v2/organizations/"
-            f"{organization}/pipelines/{pipeline}/builds"
-        )
-
-    def _request(
-        self,
-        *,
-        body: Mapping[str, Any] | None = None,
-        method: str = "GET",
-        path: str = "",
-        query: Sequence[tuple[str, str]] = (),
-    ) -> Any:
-        if not self.token:
-            raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
-        url = f"{self.base_url}{path}"
-        if query:
-            url = f"{url}?{urllib.parse.urlencode(query)}"
-        return self.transport.request(
-            url,
-            body=body,
-            headers={
-                "Authorization": f"Bearer {self.token}",
-                "Content-Type": "application/json",
-                "User-Agent": "vllm-ci-command",
-            },
-            method=method,
-        )
-
-    def list_builds(
-        self,
-        commit: str,
-        *,
-        metadata: tuple[str, str] | None = None,
-    ) -> list[dict[str, Any]]:
-        query = [
-            ("commit", commit),
-            ("exclude_jobs", "true"),
-            ("exclude_pipeline", "true"),
-            ("per_page", "100"),
-        ]
-        if metadata:
-            key, value = metadata
-            query.append((f"meta_data[{key}]", value))
-        response = self._request(query=query)
-        if not isinstance(response, list):
-            raise ApiError(None, "Buildkite API returned an invalid build list.")
-        return response
-
-    def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
-        return self._request(body=body, method="POST")
-
-    def retry_failed_jobs(
-        self,
-        build_number: int,
-        states: str,
-    ) -> dict[str, Any]:
-        number = urllib.parse.quote(str(build_number), safe="")
-        return self._request(
-            body={"states": states},
-            method="PUT",
-            path=f"/{number}/retry_failed_jobs",
-        )
-
-
-def parse_command(body: str) -> str | None:
-    if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}:
-        return body
-    return None
-
-
-def parse_trusted_users(value: str = "") -> set[str]:
-    return {
-        user.casefold() for item in value.split(",") for user in item.split() if user
-    }
-
-
-def has_ready_label(pr: Mapping[str, Any]) -> bool:
-    return any(label["name"] in READY_LABELS for label in pr["labels"])
-
-
-def is_trusted_permission(permission: str) -> bool:
-    return permission in TRUSTED_PERMISSIONS
-
-
-def authorize(
-    *,
-    actor: str,
-    permission: str,
-    pr: Mapping[str, Any],
-    trusted_approval: bool = False,
-    trusted_users: set[str] | None = None,
-) -> tuple[bool, str]:
-    trusted_users = trusted_users or set()
-    if is_trusted_permission(permission):
-        return True, f"repository {permission} permission"
-    if actor.casefold() in trusted_users:
-        return True, "configured trusted contributor"
-    if actor.casefold() != pr["user"]["login"].casefold():
-        return (
-            False,
-            "Only reviewers with write access can run CI before it is "
-            "delegated to the PR author.",
-        )
-    if pr["draft"]:
-        return False, "PR authors cannot run CI while the PR is a draft."
-    if has_ready_label(pr):
-        return True, "ready label"
-    if trusted_approval:
-        return True, "approval from a trusted reviewer"
-    return (
-        False,
-        "A reviewer with write access must run `/ci run`, approve the PR, "
-        "or add the `ready` label first.",
-    )
-
-
-def has_trusted_approval(
-    github: GitHubClient,
-    number: int,
-    trusted_users: set[str],
-) -> bool:
-    if github.get_review_decision(number) != "APPROVED":
-        return False
-
-    latest_review_states: dict[str, tuple[str, str]] = {}
-    for review in github.list_reviews(number):
-        user = review.get("user") or {}
-        login = user.get("login")
-        state = review.get("state")
-        if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
-            latest_review_states[login.casefold()] = (login, state)
-
-    for login, state in latest_review_states.values():
-        if state != "APPROVED":
-            continue
-        if login.casefold() in trusted_users:
-            return True
-        if is_trusted_permission(github.get_permission(login)):
-            return True
-    return False
-
-
-def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
-    pull_request = build.get("pull_request")
-    if isinstance(pull_request, Mapping):
-        build_pr_number = pull_request.get("id", pull_request.get("number"))
-        if build_pr_number is not None:
-            return str(build_pr_number) == str(pr_number)
-    metadata = build.get("meta_data") or {}
-    return str(metadata.get("github-pr-number")) == str(pr_number)
-
-
-def is_active_build(build: Mapping[str, Any]) -> bool:
-    return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES
-
-
-def select_latest_build(
-    builds: Sequence[dict[str, Any]],
-    pr_number: int,
-) -> dict[str, Any] | None:
-    matching = [build for build in builds if is_build_for_pr(build, pr_number)]
-    return max(matching, key=lambda build: build.get("created_at", ""), default=None)
-
-
-def create_build_payload(
-    *,
-    actor: str,
-    comment_id: int,
-    pr: Mapping[str, Any],
-) -> dict[str, Any]:
-    return {
-        "commit": pr["head"]["sha"],
-        "branch": pr["head"]["ref"],
-        "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}",
-        "pull_request_id": pr["number"],
-        "pull_request_base_branch": pr["base"]["ref"],
-        "pull_request_repository": pr["head"]["repo"]["clone_url"],
-        "pull_request_labels": [label["name"] for label in pr["labels"]],
-        "ignore_pipeline_branch_filters": True,
-        "env": {
-            "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
-            "VLLM_CI_TRIGGERED_BY": actor,
-        },
-        "meta_data": {
-            "github-comment-id": str(comment_id),
-            "github-pr-number": str(pr["number"]),
-            "github-triggered-by": actor,
-        },
-    }
-
-
-def add_reaction_safely(
-    github: GitHubClient,
-    comment_id: int,
-    content: str,
-) -> None:
-    try:
-        github.add_reaction(comment_id, content)
-    except Exception as error:
-        print(f"Could not add {content} reaction: {error}", file=sys.stderr)
-
-
-def is_already_handled(github: GitHubClient, comment_id: int) -> bool:
-    return any(
-        reaction.get("content") in {"rocket", "-1"}
-        and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
-        for reaction in github.list_reactions(comment_id)
-    )
-
-
-def handle_run_ci(
-    *,
-    actor: str,
-    buildkite: BuildkiteClient,
-    comment_id: int,
-    github: GitHubClient,
-    pr: Mapping[str, Any],
-) -> str:
-    duplicate_builds = buildkite.list_builds(
-        pr["head"]["sha"],
-        metadata=("github-comment-id", str(comment_id)),
-    )
-    duplicate = select_latest_build(duplicate_builds, pr["number"])
-    if duplicate:
-        return f"CI was already requested by this comment: {duplicate['web_url']}"
-
-    current_builds = buildkite.list_builds(pr["head"]["sha"])
-    active_build = next(
-        (
-            build
-            for build in current_builds
-            if is_build_for_pr(build, pr["number"]) and is_active_build(build)
-        ),
-        None,
-    )
-    if active_build:
-        return f"CI is already running for this commit: {active_build['web_url']}"
-
-    current_pr = github.get_pr(pr["number"])
-    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
-        return (
-            "The PR head changed while processing the command. Comment `/ci run` again."
-        )
-
-    build = buildkite.create_build(
-        create_build_payload(
-            actor=actor,
-            comment_id=comment_id,
-            pr=current_pr,
-        )
-    )
-    return (
-        f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) "
-        f"for commit `{current_pr['head']['sha'][:12]}`."
-    )
-
-
-def handle_retry_failed(
-    *,
-    buildkite: BuildkiteClient,
-    pr: Mapping[str, Any],
-) -> str:
-    builds = buildkite.list_builds(pr["head"]["sha"])
-    build = select_latest_build(builds, pr["number"])
-    if not build:
-        return "No CI build exists for the current PR commit. Use `/ci run` first."
-    if not build.get("finished_at") or is_active_build(build):
-        return f"CI is still running for this commit: {build['web_url']}"
-
-    retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
-    if retried["retried_jobs_count"] == 0:
-        return (
-            f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}"
-        )
-    return (
-        f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
-        f"[Buildkite CI #{build['number']}]({build['web_url']})."
-    )
-
-
-def run(
-    event: Mapping[str, Any],
-    github: GitHubClient,
-    buildkite: BuildkiteClient,
-    trusted_users_value: str = "",
-) -> None:
-    command = parse_command(event["comment"]["body"])
-    if not command or "pull_request" not in event["issue"]:
-        return
-
-    issue_number = event["issue"]["number"]
-    comment_id = event["comment"]["id"]
-    actor = event["comment"]["user"]["login"]
-
-    if is_already_handled(github, comment_id):
-        print(f"Comment {comment_id} was already handled.")
-        return
-    add_reaction_safely(github, comment_id, "eyes")
-
-    try:
-        pr = github.get_pr(issue_number)
-        permission = github.get_permission(actor)
-        if pr["state"] != "open":
-            github.add_comment(issue_number, "CI commands require an open PR.")
-            return
-
-        trusted_users = parse_trusted_users(trusted_users_value)
-        should_check_approval = (
-            not is_trusted_permission(permission)
-            and actor.casefold() not in trusted_users
-            and actor.casefold() == pr["user"]["login"].casefold()
-            and not pr["draft"]
-            and not has_ready_label(pr)
-        )
-        trusted_approval = should_check_approval and has_trusted_approval(
-            github,
-            issue_number,
-            trusted_users,
-        )
-        allowed, reason = authorize(
-            actor=actor,
-            permission=permission,
-            pr=pr,
-            trusted_approval=trusted_approval,
-            trusted_users=trusted_users,
-        )
-        if not allowed:
-            add_reaction_safely(github, comment_id, "-1")
-            github.add_comment(issue_number, f"@{actor}, {reason}")
-            return
-
-        print(f"Authorized @{actor}: {reason}")
-        if command == COMMAND_RUN_CI:
-            message = handle_run_ci(
-                actor=actor,
-                buildkite=buildkite,
-                comment_id=comment_id,
-                github=github,
-                pr=pr,
-            )
-        else:
-            message = handle_retry_failed(buildkite=buildkite, pr=pr)
-        add_reaction_safely(github, comment_id, "rocket")
-        github.add_comment(issue_number, message)
-    except Exception:
-        add_reaction_safely(github, comment_id, "confused")
-        raise
-
-
-def main() -> None:
-    event_path = os.environ["GITHUB_EVENT_PATH"]
-    with open(event_path, encoding="utf-8") as event_file:
-        event = json.load(event_file)
-
-    if not parse_command(event["comment"]["body"]):
-        return
-
-    github = GitHubClient(
-        os.environ.get("GH_TOKEN", ""),
-        os.environ["GITHUB_REPOSITORY"],
-    )
-    buildkite = BuildkiteClient(
-        os.environ.get("BUILDKITE_API_TOKEN", ""),
-        os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
-        os.environ.get("BUILDKITE_PIPELINE", "ci"),
-    )
-    run(
-        event,
-        github,
-        buildkite,
-        os.environ.get("CI_TRUSTED_USERS", ""),
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,363 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import unittest
-from typing import Any
-
-from run_ci_command import (
-    COMMAND_RETRY_FAILED,
-    COMMAND_RUN_CI,
-    RETRY_STATES,
-    BuildkiteClient,
-    authorize,
-    create_build_payload,
-    has_trusted_approval,
-    is_active_build,
-    is_build_for_pr,
-    parse_command,
-    parse_trusted_users,
-    run,
-    select_latest_build,
-)
-
-
-def make_pr(**overrides: Any) -> dict[str, Any]:
-    pr = {
-        "base": {"ref": "main"},
-        "draft": False,
-        "head": {
-            "ref": "feature",
-            "repo": {"clone_url": "https://github.com/contributor/vllm.git"},
-            "sha": "0123456789abcdef",
-        },
-        "labels": [],
-        "number": 42,
-        "state": "open",
-        "user": {"login": "author"},
-    }
-    pr.update(overrides)
-    return pr
-
-
-def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]:
-    return {
-        "comment": {
-            "body": command,
-            "id": 99,
-            "user": {"login": actor},
-        },
-        "issue": {
-            "number": 42,
-            "pull_request": {},
-        },
-    }
-
-
-class FakeGitHub:
-    def __init__(
-        self,
-        *,
-        permission: str = "write",
-        permissions: dict[str, str] | None = None,
-        pr: dict[str, Any] | None = None,
-        review_decision: str = "REVIEW_REQUIRED",
-        reviews: list[dict[str, Any]] | None = None,
-    ) -> None:
-        self.comments: list[str] = []
-        self.permission = permission
-        self.permissions = permissions or {}
-        self.pr = pr or make_pr()
-        self.reactions: list[str] = []
-        self.review_decision = review_decision
-        self.reviews = reviews or []
-
-    def get_pr(self, number: int) -> dict[str, Any]:
-        return self.pr
-
-    def get_permission(self, actor: str) -> str:
-        return self.permissions.get(actor, self.permission)
-
-    def get_review_decision(self, number: int) -> str:
-        return self.review_decision
-
-    def list_reviews(self, number: int) -> list[dict[str, Any]]:
-        return self.reviews
-
-    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
-        return []
-
-    def add_reaction(self, comment_id: int, content: str) -> None:
-        self.reactions.append(content)
-
-    def add_comment(self, issue_number: int, body: str) -> None:
-        self.comments.append(body)
-
-
-class FakeBuildkite:
-    def __init__(
-        self,
-        build_lists: list[list[dict[str, Any]]] | None = None,
-    ) -> None:
-        self.build_lists = build_lists or []
-        self.created_builds: list[dict[str, Any]] = []
-        self.list_calls: list[tuple[str, tuple[str, str] | None]] = []
-        self.retry_calls: list[tuple[int, str]] = []
-
-    def list_builds(
-        self,
-        commit: str,
-        *,
-        metadata: tuple[str, str] | None = None,
-    ) -> list[dict[str, Any]]:
-        self.list_calls.append((commit, metadata))
-        return self.build_lists.pop(0)
-
-    def create_build(self, body: dict[str, Any]) -> dict[str, Any]:
-        self.created_builds.append(body)
-        return {
-            "number": 123,
-            "web_url": "https://buildkite.example/builds/123",
-        }
-
-    def retry_failed_jobs(
-        self,
-        build_number: int,
-        states: str,
-    ) -> dict[str, Any]:
-        self.retry_calls.append((build_number, states))
-        return {"retried_jobs_count": 3}
-
-
-class FakeTransport:
-    def __init__(self, response: Any) -> None:
-        self.calls: list[dict[str, Any]] = []
-        self.response = response
-
-    def request(self, url: str, **kwargs: Any) -> Any:
-        self.calls.append({"url": url, **kwargs})
-        return self.response
-
-
-class RunCiCommandTest(unittest.TestCase):
-    def test_only_exact_ci_commands_are_accepted(self) -> None:
-        self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI)
-        self.assertEqual(
-            parse_command(COMMAND_RETRY_FAILED),
-            COMMAND_RETRY_FAILED,
-        )
-        self.assertIsNone(parse_command("/ci run please"))
-        self.assertIsNone(parse_command(" /ci run"))
-
-    def test_write_access_authorizes_reviewers_and_authors(self) -> None:
-        allowed, _ = authorize(
-            actor="reviewer",
-            permission="write",
-            pr=make_pr(),
-        )
-        self.assertTrue(allowed)
-
-    def test_configured_trusted_contributors_can_run_ci(self) -> None:
-        trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO")
-        allowed, _ = authorize(
-            actor="trusted-two",
-            permission="read",
-            pr=make_pr(),
-            trusted_users=trusted_users,
-        )
-        self.assertTrue(allowed)
-
-    def test_authors_need_an_approval_or_ready_label(self) -> None:
-        pending, _ = authorize(
-            actor="author",
-            permission="read",
-            pr=make_pr(),
-        )
-        approved, _ = authorize(
-            actor="author",
-            permission="read",
-            pr=make_pr(),
-            trusted_approval=True,
-        )
-        ready, _ = authorize(
-            actor="author",
-            permission="read",
-            pr=make_pr(labels=[{"name": "ready"}]),
-        )
-        self.assertFalse(pending)
-        self.assertTrue(approved)
-        self.assertTrue(ready)
-
-    def test_non_author_contributors_without_write_are_denied(self) -> None:
-        allowed, _ = authorize(
-            actor="contributor",
-            permission="read",
-            pr=make_pr(),
-            trusted_approval=True,
-        )
-        self.assertFalse(allowed)
-
-    def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None:
-        allowed, _ = authorize(
-            actor="author",
-            permission="read",
-            pr=make_pr(draft=True, labels=[{"name": "ready"}]),
-            trusted_approval=True,
-        )
-        self.assertFalse(allowed)
-
-    def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None:
-        approved_review = {
-            "state": "APPROVED",
-            "user": {"login": "reviewer"},
-        }
-        trusted = FakeGitHub(
-            permission="read",
-            permissions={"reviewer": "write"},
-            review_decision="APPROVED",
-            reviews=[approved_review],
-        )
-        untrusted = FakeGitHub(
-            permission="read",
-            review_decision="APPROVED",
-            reviews=[approved_review],
-        )
-        self.assertTrue(has_trusted_approval(trusted, 42, set()))
-        self.assertFalse(has_trusted_approval(untrusted, 42, set()))
-
-    def test_build_matching_is_scoped_to_the_pr(self) -> None:
-        self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42))
-        self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42))
-        self.assertTrue(
-            is_build_for_pr(
-                {"meta_data": {"github-pr-number": "42"}},
-                42,
-            )
-        )
-
-    def test_latest_build_selection_ignores_other_prs(self) -> None:
-        latest = select_latest_build(
-            [
-                {
-                    "created_at": "2026-07-28T02:00:00Z",
-                    "number": 3,
-                    "pull_request": {"id": 43},
-                },
-                {
-                    "created_at": "2026-07-28T01:00:00Z",
-                    "number": 2,
-                    "pull_request": {"id": 42},
-                },
-                {
-                    "created_at": "2026-07-28T00:00:00Z",
-                    "number": 1,
-                    "pull_request": {"id": 42},
-                },
-            ],
-            42,
-        )
-        self.assertEqual(latest["number"], 2)
-
-    def test_active_build_states_prevent_duplicate_runs(self) -> None:
-        self.assertTrue(is_active_build({"state": "scheduled"}))
-        self.assertTrue(is_active_build({"state": "running"}))
-        self.assertTrue(is_active_build({"state": "waiting"}))
-        self.assertTrue(is_active_build({"blocked": True, "state": "passed"}))
-        self.assertFalse(is_active_build({"state": "failed"}))
-
-    def test_build_payload_preserves_pr_context(self) -> None:
-        payload = create_build_payload(
-            actor="reviewer",
-            comment_id=99,
-            pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]),
-        )
-        self.assertEqual(
-            payload,
-            {
-                "commit": "0123456789abcdef",
-                "branch": "feature",
-                "message": "PR #42 /ci run by @reviewer",
-                "pull_request_id": 42,
-                "pull_request_base_branch": "main",
-                "pull_request_repository": ("https://github.com/contributor/vllm.git"),
-                "pull_request_labels": ["ready", "v1"],
-                "ignore_pipeline_branch_filters": True,
-                "env": {
-                    "VLLM_CI_GITHUB_COMMENT_ID": "99",
-                    "VLLM_CI_TRIGGERED_BY": "reviewer",
-                },
-                "meta_data": {
-                    "github-comment-id": "99",
-                    "github-pr-number": "42",
-                    "github-triggered-by": "reviewer",
-                },
-            },
-        )
-
-    def test_ci_run_dispatches_build_with_current_pr_metadata(self) -> None:
-        github = FakeGitHub()
-        buildkite = FakeBuildkite([[], []])
-        run(make_event(COMMAND_RUN_CI), github, buildkite)
-
-        self.assertEqual(len(buildkite.created_builds), 1)
-        self.assertEqual(
-            buildkite.created_builds[0]["message"],
-            "PR #42 /ci run by @reviewer",
-        )
-        self.assertEqual(github.reactions, ["eyes", "rocket"])
-        self.assertIn("Buildkite CI #123", github.comments[0])
-
-    def test_unapproved_authors_are_denied_without_buildkite(self) -> None:
-        github = FakeGitHub(
-            permission="read",
-            pr=make_pr(),
-            review_decision="REVIEW_REQUIRED",
-        )
-        buildkite = FakeBuildkite()
-        run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
-
-        self.assertEqual(buildkite.list_calls, [])
-        self.assertEqual(github.reactions, ["eyes", "-1"])
-        self.assertIn("approve the PR", github.comments[0])
-
-    def test_ci_retry_uses_latest_current_sha_build(self) -> None:
-        github = FakeGitHub(
-            permission="read",
-            pr=make_pr(labels=[{"name": "ready"}]),
-        )
-        buildkite = FakeBuildkite(
-            [
-                [
-                    {
-                        "created_at": "2026-07-28T01:00:00Z",
-                        "finished_at": "2026-07-28T02:00:00Z",
-                        "number": 123,
-                        "pull_request": {"id": 42},
-                        "state": "failed",
-                        "web_url": "https://buildkite.example/builds/123",
-                    }
-                ]
-            ]
-        )
-        run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
-
-        self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
-        self.assertIn("Queued 3 failed job", github.comments[0])
-
-    def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None:
-        transport = FakeTransport({"retried_jobs_count": 2})
-        client = BuildkiteClient(
-            "secret",
-            "vllm",
-            "ci",
-            transport=transport,
-        )
-        client.retry_failed_jobs(123, RETRY_STATES)
-
-        call = transport.calls[0]
-        self.assertEqual(call["method"], "PUT")
-        self.assertTrue(call["url"].endswith("/123/retry_failed_jobs"))
-        self.assertEqual(call["body"], {"states": RETRY_STATES})
-
-
-if __name__ == "__main__":
-    unittest.main()
@@ -173,6 +173,9 @@ venv.bak/

 # mkdocs documentation
 /site
+docs/argparse
+docs/examples/*
+!docs/examples/README.md

 # mypy
 .mypy_cache/
@@ -3,9 +3,6 @@ MD007:
 MD013: false
 MD024:
  siblings_only: true
-MD025:
-  # Allow front matter title to be different from the first heading in the document.
-  front_matter_title: ""
 MD031:
  list_items: false
 MD033: false
@@ -4,7 +4,7 @@ default_install_hook_types:
 default_stages:
  - pre-commit # Run locally
  - manual # Run in CI
-exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
+exclude: 'vllm/third_party/.*'
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
  rev: v0.14.0
@@ -30,7 +30,7 @@ repos:
  - id: markdownlint-cli2
    language_version: lts
    args: [--fix]
-    exclude: (^|/)CLAUDE\.md$
+    exclude: ^CLAUDE\.md$
 - repo: https://github.com/rhysd/actionlint
  rev: v1.7.7
  hooks:
@@ -260,6 +260,10 @@ repos:
    files: ^docker/(Dockerfile|versions\.json)$
    pass_filenames: false
    additional_dependencies: [dockerfile-parse]
+  - id: attention-backend-docs
+    name: Check attention backend documentation is up to date
+    entry: python tools/pre_commit/generate_attention_backend_docs.py --check
+    language: python
  - id: check-boolean-context-manager
    name: Check for boolean ops in with-statements
    entry: python tools/pre_commit/check_boolean_context_manager.py
@@ -68,8 +68,8 @@ endif()
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
-set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
 # TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
 # version check would always warn. Only treat it as a nightly build when the
 # value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
@@ -114,11 +114,6 @@ find_package(Torch REQUIRED)
 # Supported NVIDIA architectures.
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
-   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
-  # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
-  # targeting it directly.
-  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
-elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
   # to support the whole generation without specifying all sub-architectures
@@ -219,8 +214,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # the set of architectures we want to compile for and remove the from the
  # CMAKE_CUDA_FLAGS so that they are not applied globally.
  #
-  # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really
-  # needs PTX, add `+PTX` to that kernel's component-specific arch list below.
+  # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch
+  # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only
+  # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's
+  # component-specific arch list below.
  #
  clear_cuda_arches(CUDA_ARCH_FLAGS)
  extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")
@@ -230,13 +227,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  cuda_archs_loose_intersection(CUDA_ARCHS
    "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
  message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
-  if(NOT CUDA_ARCHS)
-    message(FATAL_ERROR
-      "No supported CUDA architectures; the build would produce a binary "
-      "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; "
-      "supported: ${CUDA_SUPPORTED_ARCHS}. "
-      "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).")
-  endif()
 else()
  #
  # For other GPU targets override the GPU architectures detected by cmake/torch
@@ -421,11 +411,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
-    "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
-    "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
-    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
+    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -433,7 +420,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
-        "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
    else()
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -708,7 +695,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -828,7 +815,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
  # require CUDA 12.8 or later
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -912,7 +899,7 @@ 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;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -937,7 +924,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # moe_data.cu is used by all CUTLASS MoE kernels.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -994,7 +981,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
  # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1060,7 +1047,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # Runtime dispatch is gated in
  # vllm/v1/attention/backends/mla/cutlass_mla.py.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1082,41 +1069,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set(MLA_ARCHS)
  endif()

-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
-      "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
-  endif()
-  if(FUSED_KDA_DECODE_ARCHS)
-    set(FUSED_KDA_DECODE_SRC
-      "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${FUSED_KDA_DECODE_SRC}"
-      CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
-    set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
-      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
-    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
-    message(STATUS
-      "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
-  endif()
-
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
-      "10.0f" "${CUDA_ARCHS}")
-  endif()
-  if(KIMI_K3_ATTN_RES_ARCHS)
-    set(KIMI_K3_ATTN_RES_SRC
-      "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${KIMI_K3_ATTN_RES_SRC}"
-      CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
-    set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
-      COMPILE_OPTIONS
-      "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
-    list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
-    message(STATUS
-      "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
-  endif()
-
  # Hadacore kernels
  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if(HADACORE_ARCHS)
@@ -1158,14 +1110,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_COOPERATIVE_TOPK=1)
    endif()
-    if(FUSED_KDA_DECODE_ARCHS)
-      target_compile_definitions(_C_stable_libtorch PRIVATE
-        VLLM_ENABLE_FUSED_KDA_DECODE=1)
-    endif()
-    if(KIMI_K3_ATTN_RES_ARCHS)
-      target_compile_definitions(_C_stable_libtorch PRIVATE
-        VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
-    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -1463,7 +1407,6 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/deepgemm.cmake)
    include(cmake/external_projects/fmha_sm100.cmake)
    include(cmake/external_projects/flashmla.cmake)
-    include(cmake/external_projects/flashkda.cmake)
    include(cmake/external_projects/qutlass.cmake)
    include(cmake/external_projects/tml_fa4.cmake)

@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
 - Tool calling and reasoning parsers
 - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
 - Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
+- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.

 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:

@@ -1358,10 +1358,6 @@ def main():
                            profile_memory=args.profile_memory,
                            warmup_ms=args.warmup_ms,
                            prefill_backend=pb,
-                            kv_lora_rank=args.kv_lora_rank,
-                            qk_nope_head_dim=args.qk_nope_head_dim,
-                            qk_rope_head_dim=args.qk_rope_head_dim,
-                            v_head_dim=args.v_head_dim,
                        )

                        result = run_benchmark(config)
@@ -69,11 +69,12 @@ def make_inputs(total_tokens, num_reqs, block_size):
    # Output workspace
    dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")

+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
    workspace_starts_t = torch.tensor(
        workspace_starts, dtype=torch.int32, device="cuda"
    )

-    return cache, dst, block_table, workspace_starts_t
+    return cache, dst, block_table, seq_lens_t, workspace_starts_t


 def bench_scenario(label, num_reqs, total_tokens_list, save_path):
@@ -93,7 +94,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
        )
    )
    def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, ws_starts = make_inputs(
+        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
            total_tokens, num_reqs, BLOCK_SIZE
        )

@@ -101,7 +102,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):

        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, ws_starts, num_reqs
+                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
            ),
            quantiles=quantiles,
            rep=500,
@@ -1,176 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import statistics
-
-import torch
-from tabulate import tabulate
-
-from vllm.models.inkling.nvidia.ops import qkvr_prep
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-
-def make_inputs(tokens: int, tp_size: int, is_local: bool):
-    torch.manual_seed(0)
-    num_q_heads = 64 // tp_size
-    num_kv_heads = (16 if is_local else 8) // tp_size
-    head_dim = 128
-    d_rel = 16
-    rel_extent = 512 if is_local else 1024
-    page_size = 16
-    num_blocks = (tokens + page_size - 1) // page_size
-    q_width = num_q_heads * head_dim
-    kv_width = num_kv_heads * head_dim
-    r_width = num_q_heads * d_rel
-    device = "cuda"
-
-    qkvr = torch.randn(
-        tokens,
-        q_width + 2 * kv_width + r_width,
-        device=device,
-        dtype=torch.bfloat16,
-    )
-    k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
-    v_weight = torch.randn_like(k_weight)
-    q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
-    k_norm_weight = torch.randn_like(q_norm_weight)
-    rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
-    conv_cache = torch.zeros(
-        num_blocks,
-        num_kv_heads,
-        page_size,
-        2 * head_dim,
-        device=device,
-        dtype=torch.bfloat16,
-    )
-    key_cache = torch.empty(
-        num_blocks,
-        page_size,
-        num_kv_heads,
-        head_dim,
-        device=device,
-        dtype=torch.bfloat16,
-    )
-    value_cache = torch.empty_like(key_cache)
-    positions = torch.arange(tokens, device=device, dtype=torch.int64)
-    block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
-    seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
-    slots = torch.arange(tokens, device=device, dtype=torch.int64)
-    query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
-    log_scaling = None
-    if not is_local:
-        effective_n = (positions + 1).to(torch.float32)
-        log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0))
-    return (
-        qkvr,
-        k_weight,
-        v_weight,
-        q_norm_weight,
-        k_norm_weight,
-        rel_proj,
-        1e-6,
-        num_q_heads,
-        num_kv_heads,
-        head_dim,
-        d_rel,
-        conv_cache,
-        key_cache,
-        value_cache,
-        positions,
-        block_table,
-        seq_idx,
-        slots,
-        query_start,
-        slots,
-        0,
-        head_dim,
-        page_size,
-        log_scaling,
-    )
-
-
-def capture(implementation, inputs):
-    outputs = []
-
-    def run():
-        outputs[:] = implementation.fused_qkvr_prep(*inputs)
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            run()
-    torch.cuda.current_stream().wait_stream(stream)
-    torch.accelerator.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        run()
-    torch.accelerator.synchronize()
-    return graph, outputs
-
-
-def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float:
-    for _ in range(warmup):
-        graph.replay()
-    torch.accelerator.synchronize()
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(repeats):
-        graph.replay()
-    end.record()
-    end.synchronize()
-    return start.elapsed_time(end) * 1000 / repeats
-
-
-def benchmark(inputs, args) -> float:
-    graph, _ = capture(qkvr_prep, inputs)
-    return statistics.median(
-        time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials)
-    )
-
-
-@torch.inference_mode()
-def main(args):
-    rows = []
-    for tp_size in args.tp_sizes:
-        for tokens in args.tokens:
-            for is_local in (True, False):
-                triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args)
-                rows.append(
-                    [
-                        tp_size,
-                        tokens,
-                        "local" if is_local else "global",
-                        triton_us,
-                    ]
-                )
-
-    print("Inkling QKVR prep (CUDA graph, median latency)")
-    print(
-        tabulate(
-            rows,
-            headers=[
-                "TP",
-                "tokens",
-                "scope",
-                "Triton (us)",
-            ],
-            floatfmt=("d", "d", "", ".2f"),
-        )
-    )
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    parser.add_argument(
-        "--tokens",
-        type=int,
-        nargs="+",
-        default=[1 << power for power in range(15)],
-    )
-    parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8])
-    parser.add_argument("--warmup", type=int, default=20)
-    parser.add_argument("--repeats", type=int, default=200)
-    parser.add_argument("--trials", type=int, default=5)
-    main(parser.parse_args())
@@ -1,367 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
-
-The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
-with FP32 accumulation and BF16 output. Both backends execute through CUDA
-Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
-comparison models the full latent MoE projection-and-add path.
-"""
-
-from __future__ import annotations
-
-import argparse
-import dataclasses
-import importlib.util
-import json
-import math
-import statistics
-from collections.abc import Callable, Sequence
-from pathlib import Path
-from typing import Any
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cuda.bindings import driver as cuda
-from cuda.bindings.driver import CUstream
-from quack.compile_utils import make_fake_tensor
-
-N = 7168
-K = 3584
-
-
-@dataclasses.dataclass(frozen=True, slots=True)
-class Config:
-    block_size: int
-    outputs_per_block: int
-    k_unroll: int
-    vector_width: int = 8
-
-
-def parse_config(value: str) -> Config:
-    try:
-        parts = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
-        ) from error
-    if len(parts) == 3:
-        return Config(*parts)
-    if len(parts) == 4:
-        return Config(*parts)
-    raise argparse.ArgumentTypeError(
-        "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
-    )
-
-
-def production_residual_config(m: int) -> Config | None:
-    """The measured Latent-MoE residual config for M, from the K3 table."""
-    from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
-
-    spec = KIMI_K3_PROJECTIONS.get((N, K))
-    config = spec.residual_config(m) if spec is not None else None
-    if config is None:
-        return None
-    return Config(
-        config.block_size,
-        config.outputs_per_block,
-        config.k_unroll,
-        config.vector_width,
-    )
-
-
-def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
-    if mode == "selected":
-        if selected is not None:
-            return [selected]
-        # No explicit --config: fall back to the production table for this M.
-        config = production_residual_config(m)
-        return [config] if config is not None else []
-    if mode == "baseline":
-        return [Config(224, 4, 2)]
-    return [
-        Config(block_size, outputs_per_block, k_unroll, vector_width)
-        for vector_width in (4, 8)
-        for block_size in (32, 64, 128, 224, 448)
-        if block_size % 32 == 0 and K % (block_size * vector_width) == 0
-        for outputs_per_block in (1, 2, 4, 7, 8)
-        if N % outputs_per_block == 0
-        for k_unroll in (1, 2, 4)
-    ]
-
-
-def load_kernel_class(path: Path):
-    spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
-    if spec is None or spec.loader is None:
-        raise RuntimeError(f"cannot load CuTe kernel from {path}")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module.CuteSkinnyGemm
-
-
-def stream() -> CUstream:
-    return CUstream(torch.cuda.current_stream().cuda_stream)
-
-
-def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
-    element_type = cutlass.BFloat16
-    n = cute.sym_int(divisibility=config.outputs_per_block)
-    k = cute.sym_int(divisibility=config.block_size * config.vector_width)
-    a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
-    b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
-    residual = make_fake_tensor(element_type, (m, n), divisibility=1)
-    c = make_fake_tensor(element_type, (m, n), divisibility=1)
-    kernel = kernel_class(
-        element_type=element_type,
-        num_rows=m,
-        block_size=config.block_size,
-        outputs_per_block=config.outputs_per_block,
-        vector_width=config.vector_width,
-        k_unroll=config.k_unroll,
-        has_residual=True,
-        use_pdl=True,
-    )
-    return cute.compile(
-        kernel,
-        a,
-        b,
-        residual,
-        c,
-        stream(),
-        options=(
-            "--enable-tvm-ffi --keep-cubin "
-            f"--ptxas-options -maxrregcount={max_registers} "
-            "--ptxas-options -lineinfo"
-        ),
-    )
-
-
-def resource_usage(compiled) -> dict[str, Any]:
-    executor = getattr(compiled, "_default_executor", None)
-    context = getattr(executor, "exec_context", None)
-    functions = getattr(context, "kernel_functions", None)
-    if not functions:
-        return {"resource_metrics_available": False}
-
-    def attribute(name, function) -> int:
-        error, value = cuda.cuFuncGetAttribute(name, function)
-        if error != cuda.CUresult.CUDA_SUCCESS:
-            raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
-        return int(value)
-
-    registers = [
-        attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
-        for function in functions
-    ]
-    local_bytes = [
-        attribute(
-            cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
-            function,
-        )
-        for function in functions
-    ]
-    return {
-        "resource_metrics_available": True,
-        "registers_per_thread": max(registers, default=0),
-        "spill_bytes": max(local_bytes, default=0),
-    }
-
-
-def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
-    properties = torch.cuda.get_device_properties(0)
-    bytes_per_pair = (N * K + m * N) * 2
-    target = math.ceil(multiplier * properties.L2_cache_size)
-    return max(2, min(limit, math.ceil(target / bytes_per_pair)))
-
-
-def graph_samples(
-    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
-    activation: torch.Tensor,
-    weights: Sequence[torch.Tensor],
-    residuals: Sequence[torch.Tensor],
-    repeats: int,
-    replays: int,
-) -> tuple[list[float], list[torch.Tensor]]:
-    outputs = [torch.empty_like(residual) for residual in residuals]
-    for weight, residual, output in zip(weights, residuals, outputs):
-        launch(activation, weight, residual, output)
-    torch.accelerator.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        for weight, residual, output in zip(weights, residuals, outputs):
-            launch(activation, weight, residual, output)
-    for _ in range(20):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    samples = []
-    for _ in range(repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        start.record()
-        for _ in range(replays):
-            graph.replay()
-        end.record()
-        end.synchronize()
-        samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
-    return samples, outputs
-
-
-def summarize(samples: Sequence[float]) -> dict[str, Any]:
-    ordered = sorted(samples)
-
-    def percentile(fraction: float) -> float:
-        position = fraction * (len(ordered) - 1)
-        lower = math.floor(position)
-        upper = math.ceil(position)
-        if lower == upper:
-            return ordered[lower]
-        weight = position - lower
-        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
-
-    mean = statistics.mean(samples)
-    return {
-        "median_us": statistics.median(samples),
-        "p10_us": percentile(0.1),
-        "p90_us": percentile(0.9),
-        "mean_us": mean,
-        "cv_pct": statistics.pstdev(samples) / mean * 100.0,
-        "samples_us": list(samples),
-    }
-
-
-def correctness(
-    output: torch.Tensor,
-    activation: torch.Tensor,
-    weight: torch.Tensor,
-    residual: torch.Tensor,
-) -> dict[str, Any]:
-    actual = output.float()
-    reference = activation.float() @ weight.float().t() + residual.float()
-    error = (actual - reference).abs()
-    scaled_error = error / (reference.abs() + 1.0)
-    cosine = torch.nn.functional.cosine_similarity(
-        actual.flatten(), reference.flatten(), dim=0
-    ).item()
-    return {
-        "valid": cosine > 0.999,
-        "cosine": cosine,
-        "max_abs_error": error.max().item(),
-        "max_scaled_error": scaled_error.max().item(),
-    }
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--kernel", type=Path, required=True)
-    parser.add_argument("--output", type=Path, required=True)
-    parser.add_argument(
-        "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
-    )
-    parser.add_argument("--config", type=parse_config)
-    parser.add_argument("--m", type=int, action="append")
-    parser.add_argument("--config-shard", type=int, default=0)
-    parser.add_argument("--num-config-shards", type=int, default=1)
-    parser.add_argument("--repeats", type=int, default=21)
-    parser.add_argument("--replays", type=int, default=200)
-    parser.add_argument("--cache-multiplier", type=float, default=3.0)
-    parser.add_argument("--max-buffers", type=int, default=32)
-    parser.add_argument("--max-registers", type=int, default=64)
-    args = parser.parse_args()
-
-    token_counts = args.m or list(range(1, 17))
-    if any(not 1 <= m <= 16 for m in token_counts):
-        raise ValueError("expected 1 <= M <= 16")
-    if not 0 <= args.config_shard < args.num_config_shards:
-        raise ValueError("config shard must be in [0, num_config_shards)")
-    torch.accelerator.set_device_index(0)
-    if torch.cuda.get_device_capability() != (10, 3):
-        raise RuntimeError("this benchmark requires SM103")
-
-    kernel_class = load_kernel_class(args.kernel)
-    properties = torch.cuda.get_device_properties(0)
-    metadata = {
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability()),
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-    }
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    with args.output.open("w", encoding="utf-8") as output_file:
-        for m in token_counts:
-            configs = candidate_configs(args.mode, args.config, m)
-            torch.manual_seed(20260722 + m)
-            count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
-            activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
-            weights = [
-                torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
-                for _ in range(count)
-            ]
-            residuals = [
-                torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
-                for _ in range(count)
-            ]
-            candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
-            candidates.extend(
-                ("cute_residual", config)
-                for index, config in enumerate(configs)
-                if index % args.num_config_shards == args.config_shard
-            )
-            for backend, config in candidates:
-                row: dict[str, Any] = {
-                    "m": m,
-                    "n": N,
-                    "k": K,
-                    "backend": backend,
-                    "mode": args.mode,
-                    "config": dataclasses.asdict(config) if config else {},
-                    "num_buffers": count,
-                    "cache_multiplier": args.cache_multiplier,
-                    **metadata,
-                }
-                try:
-                    if backend == "cublas_addmm":
-                        launch = lambda a, b, residual, c: torch.addmm(
-                            residual, a, b.t(), out=c
-                        )
-                    else:
-                        if config is None:
-                            raise AssertionError("missing CuTe config")
-                        compiled = compile_kernel(
-                            kernel_class, m, config, args.max_registers
-                        )
-                        launch = lambda a, b, residual, c, fn=compiled: fn(
-                            a, b, residual, c, stream()
-                        )
-                        row.update(resource_usage(compiled))
-                    samples, outputs = graph_samples(
-                        launch,
-                        activation,
-                        weights,
-                        residuals,
-                        args.repeats,
-                        args.replays,
-                    )
-                    row.update(
-                        correctness(outputs[0], activation, weights[0], residuals[0])
-                    )
-                    row.update(summarize(samples))
-                except Exception as error:  # noqa: BLE001
-                    row.update(
-                        {
-                            "valid": False,
-                            "error": f"{type(error).__name__}: {error}",
-                        }
-                    )
-                output_file.write(json.dumps(row, sort_keys=True) + "\n")
-                output_file.flush()
-                print(json.dumps(row, sort_keys=True), flush=True)
-
-            del activation, weights, residuals
-            torch.accelerator.empty_cache()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,806 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels.
-
-The ``up-projection`` subcommand isolates the TP-local dynamic and static-M
-skinny GEMMs. It rotates weights through a working set larger than L2 to model
-successive model layers.
-
-The ``whole-tail`` subcommand measures the distributed operator. Its reference
-path includes two AllReduces, RMSNorm, the replicated up-projection, and the
-final add. CUDA-event samples report the slowest rank so cross-rank skew is
-included.
-
-Examples:
-
-.. code-block:: console
-
-    .venv/bin/python \
-      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection
-
-    torchrun --nproc-per-node=8 \
-      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
-
-For multi-node runs, launch one ``torchrun`` agent per node and use a shared
-rendezvous endpoint.
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import math
-import os
-import statistics
-from collections.abc import Callable, Sequence
-from dataclasses import asdict
-from pathlib import Path
-from typing import Any
-
-import cutlass
-import cutlass.utils as utils
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-from cuda.bindings import driver as cuda
-
-from vllm.distributed import get_tp_group
-from vllm.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-    set_custom_all_reduce,
-)
-from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
-from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
-from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
-    fused_add_multicast_gemm,
-    fused_add_multicast_skinny_gemm,
-)
-
-HIDDEN_SIZE = 7168
-LATENT_SIZE = 3584
-RMS_EPS = 0.1
-MAX_NUM_TOKENS = 16
-MMA_TILER_MN = (64, 32)
-CLUSTER_SHAPE_MN = (1, 8)
-B_PRIME_STAGES = 2
-
-
-def parse_up_projection_config(
-    value: str,
-) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
-    try:
-        values = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
-        ) from error
-    if len(values) in (3, 4):
-        return fused_add_multicast_skinny_gemm.SkinnyConfig(*values)
-    if len(values) == 5 and values[4] in (0, 1):
-        return fused_add_multicast_skinny_gemm.SkinnyConfig(
-            *values[:4],
-            prefetch_b_before_pdl=bool(values[4]),
-        )
-    raise argparse.ArgumentTypeError(
-        "config must be BLOCK,OUTPUTS,K_UNROLL"
-        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
-    )
-
-
-def parse_tail_skinny_config(
-    value: str,
-) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]:
-    try:
-        values = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
-        ) from error
-    if len(values) == 4:
-        num_tokens, *config = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
-    if len(values) == 5:
-        num_tokens, *config = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
-    if len(values) == 6 and values[5] in (0, 1):
-        num_tokens, block, outputs, unroll, vector_width, prefetch = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(
-            block,
-            outputs,
-            unroll,
-            vector_width,
-            bool(prefetch),
-        )
-    raise argparse.ArgumentTypeError(
-        "config must be M,BLOCK,OUTPUTS,K_UNROLL"
-        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
-    )
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    subparsers = parser.add_subparsers(dest="scope", required=True)
-
-    up_projection = subparsers.add_parser(
-        "up-projection",
-        help="Benchmark the isolated TP-local up-projection kernels.",
-    )
-    up_projection.add_argument(
-        "--backend",
-        choices=("dynamic", "skinny", "both"),
-        default="both",
-    )
-    up_projection.add_argument("--tp-size", type=int, default=16)
-    up_projection.add_argument(
-        "--num-tokens",
-        type=int,
-        nargs="+",
-        default=[*range(1, 9), 16],
-    )
-    up_projection.add_argument(
-        "--skinny-config",
-        type=parse_up_projection_config,
-        action="append",
-        help="Benchmark a static-M config for every selected token count.",
-    )
-    up_projection.add_argument("--cache-multiplier", type=float, default=2.0)
-    up_projection.add_argument("--max-weights", type=int, default=64)
-    up_projection.add_argument("--warmup-replays", type=int, default=10)
-    up_projection.add_argument("--samples", type=int, default=31)
-    up_projection.add_argument("--output", type=Path)
-
-    whole_tail = subparsers.add_parser(
-        "whole-tail",
-        help="Benchmark the distributed latent-MoE tail operator.",
-    )
-    whole_tail.add_argument(
-        "--backend",
-        choices=("reference", "fused", "both"),
-        default="both",
-    )
-    whole_tail.add_argument(
-        "--num-tokens",
-        type=int,
-        nargs="+",
-        default=[1, 5, 8, 16],
-    )
-    whole_tail.add_argument("--warmup-replays", type=int, default=20)
-    whole_tail.add_argument("--samples", type=int, default=51)
-    whole_tail.add_argument(
-        "--skinny-max-num-tokens",
-        type=int,
-        nargs="+",
-        help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.",
-    )
-    whole_tail.add_argument(
-        "--skinny-config",
-        type=parse_tail_skinny_config,
-        action="append",
-        help="Override one static-M config for tuning.",
-    )
-    whole_tail.add_argument("--output", type=Path)
-    return parser.parse_args()
-
-
-def percentile(samples: Sequence[float], fraction: float) -> float:
-    ordered = sorted(samples)
-    position = fraction * (len(ordered) - 1)
-    lower = math.floor(position)
-    upper = math.ceil(position)
-    if lower == upper:
-        return ordered[lower]
-    upper_weight = position - lower
-    return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight
-
-
-def summarize(samples_us: Sequence[float]) -> dict[str, Any]:
-    mean_us = statistics.mean(samples_us)
-    return {
-        "median_us": statistics.median(samples_us),
-        "p10_us": percentile(samples_us, 0.1),
-        "p90_us": percentile(samples_us, 0.9),
-        "mean_us": mean_us,
-        "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0,
-        "samples_us": list(samples_us),
-    }
-
-
-def rotating_weight_count(
-    shard_size: int,
-    cache_multiplier: float,
-    limit: int,
-) -> int:
-    properties = torch.cuda.get_device_properties(
-        torch.accelerator.current_device_index()
-    )
-    weight_bytes = shard_size * LATENT_SIZE * 2
-    target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier)
-    return max(2, min(limit, math.ceil(target_bytes / weight_bytes)))
-
-
-def capture_up_projection_graph(
-    launches: Sequence[Callable[[], None]],
-) -> torch.cuda.CUDAGraph:
-    for launch in launches:
-        launch()
-    torch.accelerator.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        for launch in launches:
-            launch()
-    torch.accelerator.synchronize()
-    return graph
-
-
-def benchmark_up_projection_graph(
-    graph: torch.cuda.CUDAGraph,
-    *,
-    operations_per_replay: int,
-    warmup_replays: int,
-    samples: int,
-) -> dict[str, Any]:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    samples_us = []
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    for _ in range(samples):
-        start.record()
-        graph.replay()
-        end.record()
-        end.synchronize()
-        samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay)
-    return summarize(samples_us)
-
-
-class DynamicKernel:
-    def __init__(
-        self,
-        shard_size: int,
-        mailbox: torch.Tensor,
-        shared_shard: torch.Tensor,
-    ) -> None:
-        self.shard_size = shard_size
-        self.mailbox = mailbox
-        self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox)
-        compile_latent = torch.empty(
-            (1, MAX_NUM_TOKENS, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=mailbox.device,
-        )
-        compile_weight = torch.empty(
-            (1, shard_size, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=mailbox.device,
-        )
-        cluster_size = math.prod(CLUSTER_SHAPE_MN)
-        max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
-        self.compiled = fused_add_multicast_gemm.compile_kernel(
-            (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1),
-            fused_add_multicast_gemm._as_cute(
-                compile_latent,
-                dynamic_m=True,
-            ),
-            fused_add_multicast_gemm._as_cute(compile_weight),
-            self.mailbox_c,
-            fused_add_multicast_gemm._as_cute(shared_shard),
-            HIDDEN_SIZE,
-            shard_size,
-            MMA_TILER_MN,
-            CLUSTER_SHAPE_MN,
-            max_active_clusters,
-            B_PRIME_STAGES,
-        )
-
-    def launch(
-        self,
-        latent: torch.Tensor,
-        weight: torch.Tensor,
-        shared_shard: torch.Tensor,
-    ) -> None:
-        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-        self.compiled(
-            fused_add_multicast_gemm._as_cute(
-                latent.unsqueeze(0),
-                dynamic_m=True,
-            ),
-            fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)),
-            self.mailbox_c,
-            fused_add_multicast_gemm._as_cute(shared_shard),
-            cutlass.Int64(latent.shape[0]),
-            cutlass.Int64(self.mailbox.data_ptr()),
-            stream,
-        )
-
-
-class SkinnyKernel:
-    def __init__(
-        self,
-        num_tokens: int,
-        shard_size: int,
-        config: fused_add_multicast_skinny_gemm.SkinnyConfig,
-    ) -> None:
-        self.compiled = fused_add_multicast_skinny_gemm.compile_kernel(
-            num_rows=num_tokens,
-            latent_dim=LATENT_SIZE,
-            hidden_dim=HIDDEN_SIZE,
-            shard_dim=shard_size,
-            config=config,
-        )
-
-    def launch(
-        self,
-        latent: torch.Tensor,
-        weight: torch.Tensor,
-        shared_shard: torch.Tensor,
-        mailbox: torch.Tensor,
-    ) -> None:
-        self.compiled(
-            fused_add_multicast_skinny_gemm._as_cute(latent),
-            fused_add_multicast_skinny_gemm._as_cute(weight),
-            fused_add_multicast_skinny_gemm._as_cute(shared_shard),
-            cutlass.Int64(mailbox.data_ptr()),
-            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
-        )
-
-
-def check_up_projection_output(
-    actual: torch.Tensor,
-    latent: torch.Tensor,
-    weight: torch.Tensor,
-    shared_shard: torch.Tensor,
-) -> None:
-    gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16)
-    expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16)
-    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
-
-
-def make_up_projection_launches(
-    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None],
-    latent: torch.Tensor,
-    weights: Sequence[torch.Tensor],
-    shared_shard: torch.Tensor,
-) -> list[Callable[[], None]]:
-    return [
-        lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights
-    ]
-
-
-def benchmark_up_projection(args: argparse.Namespace) -> None:
-    if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size:
-        raise ValueError("TP size must be positive and divide the hidden size")
-    if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens):
-        raise ValueError("--num-tokens values must be in [1, 16]")
-    if args.cache_multiplier <= 0 or args.max_weights <= 0:
-        raise ValueError("cache multiplier and max weights must be positive")
-    if args.warmup_replays < 0 or args.samples <= 0:
-        raise ValueError("warmup replays must be nonnegative and samples positive")
-
-    torch.accelerator.set_device_index(0)
-    device = torch.device("cuda", 0)
-    if torch.cuda.get_device_capability(device)[0] != 10:
-        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
-
-    shard_size = HIDDEN_SIZE // args.tp_size
-    weight_count = rotating_weight_count(
-        shard_size,
-        args.cache_multiplier,
-        args.max_weights,
-    )
-    torch.manual_seed(20260726)
-    weights = [
-        torch.randn(
-            (shard_size, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        / LATENT_SIZE**0.5
-        for _ in range(weight_count)
-    ]
-    mailbox = torch.empty(
-        (1, MAX_NUM_TOKENS, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    shared = torch.randn(
-        (MAX_NUM_TOKENS, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    shared_shard = shared[:, :shard_size]
-    use_dynamic = args.backend in ("dynamic", "both")
-    use_skinny = args.backend in ("skinny", "both")
-    dynamic_kernel = (
-        DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None
-    )
-
-    results = []
-    for num_tokens in args.num_tokens:
-        latent = torch.randn(
-            (num_tokens, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        result: dict[str, Any] = {"num_tokens": num_tokens}
-        if dynamic_kernel is not None:
-            launches = make_up_projection_launches(
-                dynamic_kernel.launch,
-                latent,
-                weights,
-                shared_shard,
-            )
-            graph = capture_up_projection_graph(launches)
-            result["dynamic"] = benchmark_up_projection_graph(
-                graph,
-                operations_per_replay=len(launches),
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-            )
-            check_up_projection_output(
-                mailbox[0, :num_tokens, :shard_size],
-                latent,
-                weights[-1],
-                shared_shard[:num_tokens],
-            )
-        if use_skinny:
-            configs = args.skinny_config or [
-                fused_add_multicast_skinny_gemm.config_for_m(
-                    num_tokens,
-                    shard_size,
-                )
-            ]
-            skinny_results = []
-            for config in configs:
-                skinny_kernel = SkinnyKernel(num_tokens, shard_size, config)
-
-                def launch_skinny(
-                    latent: torch.Tensor,
-                    weight: torch.Tensor,
-                    shared_shard: torch.Tensor,
-                    *,
-                    skinny_kernel: SkinnyKernel = skinny_kernel,
-                    num_tokens: int = num_tokens,
-                ) -> None:
-                    skinny_kernel.launch(
-                        latent,
-                        weight,
-                        shared_shard[:num_tokens],
-                        mailbox,
-                    )
-
-                launches = make_up_projection_launches(
-                    launch_skinny,
-                    latent,
-                    weights,
-                    shared_shard,
-                )
-                graph = capture_up_projection_graph(launches)
-                timing = benchmark_up_projection_graph(
-                    graph,
-                    operations_per_replay=len(launches),
-                    warmup_replays=args.warmup_replays,
-                    samples=args.samples,
-                )
-                check_up_projection_output(
-                    mailbox[0, :num_tokens, :shard_size],
-                    latent,
-                    weights[-1],
-                    shared_shard[:num_tokens],
-                )
-                skinny_results.append(
-                    {
-                        "config": asdict(config),
-                        **timing,
-                    }
-                )
-            result["skinny"] = skinny_results
-        results.append(result)
-
-    properties = torch.cuda.get_device_properties(device)
-    report = {
-        "scope": "up-projection",
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability(device)),
-        "tp_size": args.tp_size,
-        "shard_size": shard_size,
-        "weight_count": weight_count,
-        "cache_multiplier": args.cache_multiplier,
-        "warmup_replays": args.warmup_replays,
-        "samples": args.samples,
-        "results": results,
-    }
-    rendered = json.dumps(report, indent=2)
-    print(rendered, flush=True)
-    if args.output is not None:
-        args.output.parent.mkdir(parents=True, exist_ok=True)
-        args.output.write_text(rendered + "\n", encoding="utf-8")
-
-
-def capture_tail_graph(
-    operation: Callable[[], torch.Tensor],
-    cpu_group: dist.ProcessGroup,
-) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
-    for _ in range(3):
-        dist.barrier(group=cpu_group)
-        output = operation()
-    torch.accelerator.synchronize()
-
-    dist.barrier(group=cpu_group)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        output = operation()
-    torch.accelerator.synchronize()
-    return graph, output
-
-
-def benchmark_tail_graph(
-    graph: torch.cuda.CUDAGraph,
-    *,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> dict[str, Any]:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    dist.barrier(group=cpu_group)
-    starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
-    ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
-    for start, end in zip(starts, ends):
-        start.record()
-        graph.replay()
-        end.record()
-    torch.accelerator.synchronize()
-
-    samples_us = torch.tensor(
-        [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)],
-        dtype=torch.float64,
-        device=torch.accelerator.current_device_index(),
-    )
-    dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group)
-    return summarize(samples_us[1:].tolist())
-
-
-def make_inputs(
-    num_tokens: int,
-    rank: int,
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    torch.manual_seed(20260726 + 100 * num_tokens + rank)
-    routed = torch.randn(
-        (num_tokens, LATENT_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    ).mul_(0.01)
-    shared = torch.randn(
-        (num_tokens, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    return routed, shared
-
-
-def make_reference(
-    routed: torch.Tensor,
-    shared: torch.Tensor,
-    rms_weight: torch.Tensor,
-    up_weight: torch.Tensor,
-    device_group: dist.ProcessGroup,
-) -> Callable[[], torch.Tensor]:
-    routed_workspace = torch.empty_like(routed)
-    shared_workspace = torch.empty_like(shared)
-
-    def reference() -> torch.Tensor:
-        routed_workspace.copy_(routed)
-        dist.all_reduce(routed_workspace, group=device_group)
-        normalized = F.rms_norm(
-            routed_workspace,
-            (LATENT_SIZE,),
-            rms_weight,
-            RMS_EPS,
-        )
-        projected = F.linear(normalized, up_weight)
-        shared_workspace.copy_(shared)
-        dist.all_reduce(shared_workspace, group=device_group)
-        return projected.add(shared_workspace)
-
-    return reference
-
-
-def check_fused_output(
-    fused_output: torch.Tensor,
-    reference: Callable[[], torch.Tensor],
-    cpu_group: dist.ProcessGroup,
-) -> None:
-    dist.barrier(group=cpu_group)
-    expected = reference()
-    torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2)
-
-
-def benchmark_whole_tail(args: argparse.Namespace) -> None:
-    if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens):
-        raise ValueError("--num-tokens values must be in [1, 16]")
-    if args.warmup_replays < 0 or args.samples <= 0:
-        raise ValueError("warmup replays must be nonnegative and samples positive")
-    if args.skinny_max_num_tokens is not None and any(
-        not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
-    ):
-        raise ValueError("--skinny-max-num-tokens must be in [0, 8]")
-    skinny_configs = dict(args.skinny_config or ())
-    if len(skinny_configs) != len(args.skinny_config or ()):
-        raise ValueError("--skinny-config must not repeat an M value")
-    if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs):
-        raise ValueError("--skinny-config M values must be in [1, 8]")
-    if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys():
-        raise RuntimeError("launch this benchmark with torchrun")
-
-    rank = int(os.environ["RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    local_rank = int(os.environ["LOCAL_RANK"])
-    device = torch.device("cuda", local_rank)
-    torch.accelerator.set_device_index(device)
-    init_distributed_environment()
-    if world_size > 8:
-        set_custom_all_reduce(False)
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
-    device_group = get_tp_group().device_group
-    cpu_group = dist.new_group(backend="gloo")
-
-    if torch.cuda.get_device_capability(device)[0] != 10:
-        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
-
-    torch.manual_seed(20260726)
-    rms_weight = 1 + 0.1 * torch.randn(
-        LATENT_SIZE,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    up_weight = (
-        torch.randn(
-            (HIDDEN_SIZE, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        / LATENT_SIZE**0.5
-    )
-
-    use_reference = args.backend in ("reference", "both")
-    use_fused = args.backend in ("fused", "both")
-    fused_ops = []
-    if use_fused:
-        production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
-
-        def config_for_m(
-            num_rows: int,
-            shard_dim: int = 896,
-        ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
-            config = skinny_configs.get(num_rows)
-            if config is not None:
-                return config
-            return production_config_for_m(num_rows, shard_dim)
-
-        fused_add_multicast_skinny_gemm.config_for_m = config_for_m
-        cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS]
-        for cutoff in cutoffs:
-            latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
-            latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
-            fused_ops.append(
-                (
-                    cutoff,
-                    latent_moe_tail.KimiK3LatentMoETailOp.initialize(
-                        hidden_size=HIDDEN_SIZE,
-                        latent_size=LATENT_SIZE,
-                        dtype=torch.bfloat16,
-                        device=device,
-                        rms_eps=RMS_EPS,
-                    ),
-                )
-            )
-        cutedsl_warmup()
-
-    results = []
-    for num_tokens in args.num_tokens:
-        routed, shared = make_inputs(num_tokens, rank, device)
-        reference = make_reference(
-            routed,
-            shared,
-            rms_weight,
-            up_weight,
-            device_group,
-        )
-        result: dict[str, Any] = {"num_tokens": num_tokens}
-        if use_reference:
-            reference_graph, _ = capture_tail_graph(reference, cpu_group)
-            result["reference"] = benchmark_tail_graph(
-                reference_graph,
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-                device_group=device_group,
-                cpu_group=cpu_group,
-            )
-        for cutoff, fused_op in fused_ops:
-
-            def fused(
-                routed: torch.Tensor = routed,
-                shared: torch.Tensor = shared,
-                fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
-            ) -> torch.Tensor:
-                return fused_op(routed, shared, rms_weight, up_weight)
-
-            fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
-            fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
-            result[fused_key] = benchmark_tail_graph(
-                fused_graph,
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-                device_group=device_group,
-                cpu_group=cpu_group,
-            )
-            check_fused_output(fused_output, reference, cpu_group)
-            if "reference" in result:
-                speedup = (
-                    result["reference"]["median_us"] / result[fused_key]["median_us"]
-                )
-                if len(fused_ops) == 1:
-                    result["speedup"] = speedup
-                else:
-                    result[f"{fused_key}_speedup"] = speedup
-        results.append(result)
-
-    properties = torch.cuda.get_device_properties(device)
-    report = {
-        "scope": "whole-tail",
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability(device)),
-        "world_size": world_size,
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-        "warmup_replays": args.warmup_replays,
-        "samples": args.samples,
-        "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
-        "skinny_configs": {
-            str(num_tokens): asdict(config)
-            for num_tokens, config in skinny_configs.items()
-        },
-        "timing_scope": {
-            "reference": (
-                "two input copies, two AllReduces, RMSNorm, full replicated "
-                "up-projection GEMM, and final add"
-            ),
-            "fused": (
-                "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded "
-                "up-projection/multicast, and Lamport copy"
-            ),
-        },
-        "results": results,
-    }
-    if rank == 0:
-        rendered = json.dumps(report, indent=2)
-        print(rendered, flush=True)
-        if args.output is not None:
-            args.output.parent.mkdir(parents=True, exist_ok=True)
-            args.output.write_text(rendered + "\n", encoding="utf-8")
-
-    dist.barrier(group=cpu_group)
-
-
-def main() -> None:
-    args = parse_args()
-    if args.scope == "up-projection":
-        benchmark_up_projection(args)
-        return
-
-    from vllm.config import VllmConfig, set_current_vllm_config
-
-    with set_current_vllm_config(VllmConfig()):
-        benchmark_whole_tail(args)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,239 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import json
-import os
-import statistics
-from collections.abc import Callable
-
-import torch
-import torch.distributed as dist
-
-import vllm._custom_ops as ops
-from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024])
-    parser.add_argument("--hidden-size", type=int, default=7168)
-    parser.add_argument("--graph-repeats", type=int, default=20)
-    parser.add_argument("--warmup-replays", type=int, default=5)
-    parser.add_argument("--samples", type=int, default=15)
-    return parser.parse_args()
-
-
-def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph:
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            op()
-    stream.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        for _ in range(repeats):
-            op()
-    torch.cuda.current_stream().wait_stream(stream)
-    return graph
-
-
-def max_rank_graph_time(
-    graph: torch.cuda.CUDAGraph,
-    repeats: int,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> float:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    timings = []
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    for _ in range(samples):
-        dist.barrier(group=cpu_group)
-        start.record()
-        graph.replay()
-        end.record()
-        end.synchronize()
-        elapsed = torch.tensor(
-            start.elapsed_time(end) / repeats,
-            dtype=torch.float64,
-            device=torch.accelerator.current_device_index(),
-        )
-        dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
-        timings.append(elapsed.item())
-    return statistics.median(timings)
-
-
-def check_outputs(
-    comm: CustomAllreduce,
-    local: torch.Tensor,
-    reduce_input: torch.Tensor,
-    device_group: dist.ProcessGroup,
-) -> None:
-    expected_gather = torch.empty(
-        (local.shape[0] * dist.get_world_size(), local.shape[1]),
-        dtype=local.dtype,
-        device=local.device,
-    )
-    dist.all_gather_into_tensor(expected_gather, local, group=device_group)
-    gathered = comm.custom_all_gather(local)
-    assert gathered is not None
-    torch.testing.assert_close(gathered, expected_gather)
-
-    expected_scatter = torch.empty_like(local)
-    dist.reduce_scatter_tensor(
-        expected_scatter,
-        reduce_input.clone(),
-        group=device_group,
-    )
-    scattered = comm.custom_reduce_scatter(reduce_input)
-    assert scattered is not None
-    torch.testing.assert_close(scattered, expected_scatter)
-
-
-def benchmark_shape(
-    comm: CustomAllreduce,
-    global_tokens: int,
-    hidden_size: int,
-    graph_repeats: int,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> dict[str, float | int]:
-    world_size = dist.get_world_size()
-    rank = dist.get_rank()
-    padded_tokens = (global_tokens + world_size - 1) // world_size * world_size
-    local_tokens = padded_tokens // world_size
-    local = torch.full(
-        (local_tokens, hidden_size),
-        rank + 1,
-        dtype=torch.bfloat16,
-        device=torch.accelerator.current_device_index(),
-    )
-    reduce_input = torch.full(
-        (padded_tokens, hidden_size),
-        rank + 1,
-        dtype=torch.bfloat16,
-        device=local.device,
-    )
-    check_outputs(comm, local, reduce_input, device_group)
-
-    custom_gather_out = torch.empty(
-        (padded_tokens, hidden_size),
-        dtype=local.dtype,
-        device=local.device,
-    )
-    custom_scatter_out = torch.empty_like(local)
-    nccl_gather_out = torch.empty_like(custom_gather_out)
-    nccl_scatter_out = torch.empty_like(local)
-
-    def custom_ag() -> None:
-        ops.mnnvl_lamport_all_gather(
-            comm._ptr,
-            local,
-            custom_gather_out,
-            comm.mnnvl_lamport_ag_local_ptr,
-            comm.mnnvl_lamport_ag_multicast_ptr,
-            comm.mnnvl_lamport_ag_epoch_ptr,
-            comm.mnnvl_buffer_size,
-        )
-
-    def custom_rs() -> None:
-        ops.mnnvl_lamport_reduce_scatter(
-            comm._ptr,
-            reduce_input,
-            custom_scatter_out,
-            comm.mnnvl_lamport_rs_local_ptr,
-            comm.mnnvl_lamport_rs_epoch_ptr,
-            comm.mnnvl_buffer_size,
-        )
-
-    def nccl_ag() -> None:
-        dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group)
-
-    def nccl_rs() -> None:
-        dist.reduce_scatter_tensor(
-            nccl_scatter_out,
-            reduce_input,
-            group=device_group,
-        )
-
-    graphs = {
-        "custom_ag_us": capture_graph(custom_ag, graph_repeats),
-        "nccl_ag_us": capture_graph(nccl_ag, graph_repeats),
-        "custom_rs_us": capture_graph(custom_rs, graph_repeats),
-        "nccl_rs_us": capture_graph(nccl_rs, graph_repeats),
-    }
-    times = {
-        name: max_rank_graph_time(
-            graph,
-            graph_repeats,
-            warmup_replays,
-            samples,
-            device_group,
-            cpu_group,
-        )
-        * 1000
-        for name, graph in graphs.items()
-    }
-    torch.testing.assert_close(custom_gather_out, nccl_gather_out)
-    torch.testing.assert_close(custom_scatter_out, nccl_scatter_out)
-    return {
-        "global_tokens": global_tokens,
-        "padded_tokens": padded_tokens,
-        "local_bytes": local.nbytes,
-        "full_bytes": reduce_input.nbytes,
-        **times,
-        "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"],
-        "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"],
-    }
-
-
-def main() -> None:
-    args = parse_args()
-    local_rank = int(os.environ["LOCAL_RANK"])
-    torch.accelerator.set_device_index(local_rank)
-    dist.init_process_group("nccl")
-    device_group = dist.group.WORLD
-    cpu_group = dist.new_group(backend="gloo")
-
-    comm = CustomAllreduce(
-        group=cpu_group,
-        device=torch.device("cuda", local_rank),
-    )
-    assert not comm.disabled
-    assert comm.world_size == 16
-    assert comm.mnnvl_only
-    assert comm.mnnvl_multicast_ptr
-
-    results = [
-        benchmark_shape(
-            comm,
-            tokens,
-            args.hidden_size,
-            args.graph_repeats,
-            args.warmup_replays,
-            args.samples,
-            device_group,
-            cpu_group,
-        )
-        for tokens in args.tokens
-    ]
-    if dist.get_rank() == 0:
-        print(json.dumps(results, indent=2), flush=True)
-
-    comm.close()
-    dist.destroy_process_group(cpu_group)
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()
@@ -154,7 +154,7 @@ def main(
                scale=scale,
                causal=True,
                alibi_slopes=None,
-                sliding_window=window_size if sliding_window is not None else -1,
+                sliding_window=window_size,
                block_table=block_tables,
                softcap=0,
                scheduler_metadata=metadata,
@@ -1,267 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
-
-Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
-long greedy decode (CUDA graphs on) once with the standard kernel and once with
-ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
-separate subprocesses so each gets a clean CUDA context.
-
-The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
-CUDA-graph capture on the pre-release Blackwell FP4 path); pass
--no-disable-flashinfer-autotune for non-FP4 models.
-
-Examples:
-    python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
-    python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
-        --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
-"""
-
-import argparse
-import json
-import os
-import subprocess
-import sys
-import time
-
-DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
-
-MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
-
-
-def parse_args():
-    p = argparse.ArgumentParser(
-        description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
-    )
-    p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
-    p.add_argument("--prompt", default=DEFAULT_PROMPT)
-    p.add_argument("--batch-size", type=int, default=256)
-    p.add_argument("--num-steps", type=int, default=1000)
-    p.add_argument("--warmup-steps", type=int, default=128)
-    p.add_argument("--repeats", type=int, default=1)
-    p.add_argument(
-        "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
-    )
-    p.add_argument(
-        "--dtype",
-        default="bfloat16",
-        choices=["bfloat16", "float16", "float32", "auto"],
-    )
-    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    p.add_argument("--max-model-len", type=int, default=None)
-    p.add_argument(
-        "--disable-flashinfer-autotune",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
-        "It is unstable under CUDA-graph capture on the "
-        "pre-release Blackwell FP4 path; pass "
-        "--no-disable-flashinfer-autotune for non-FP4 models.",
-    )
-    p.add_argument(
-        "--mamba-ssm-cache-dtype",
-        default="auto",
-        choices=["auto", "float32", "float16", "bfloat16"],
-        help="SSM state dtype (both modes). 'auto' = config-driven; "
-        "'float32' = fp32 state, 'bfloat16' = s16 state.",
-    )
-    p.add_argument(
-        "--baseline-ssm-config",
-        default="",
-        help="Pin the STANDARD baseline's SSM launch config as "
-        "'bsm,nw' via override_ssm_config (forces the in-process "
-        "engine so the override reaches the kernel). Empty = off.",
-    )
-    p.add_argument(
-        "--worker",
-        choices=["standard", "replayssm"],
-        default=None,
-        help=argparse.SUPPRESS,
-    )
-    return p.parse_args()
-
-
-def resolve_max_model_len(args) -> int:
-    if args.max_model_len is not None:
-        return args.max_model_len
-    return args.num_steps + 256
-
-
-def run_worker(args):
-    # override_ssm_config is a module global; it only reaches the model if the
-    # engine runs in-process (default V1 spawns a separate EngineCore). Force it.
-    if args.worker == "standard" and args.baseline_ssm_config:
-        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-
-    import torch
-
-    from vllm import LLM, SamplingParams
-
-    mode = args.worker
-    max_model_len = resolve_max_model_len(args)
-
-    llm_kwargs = dict(
-        model=args.model_id,
-        tensor_parallel_size=1,
-        dtype=args.dtype,
-        max_model_len=max_model_len,
-        trust_remote_code=True,
-        enable_prefix_caching=False,
-        enable_chunked_prefill=False,
-        max_num_seqs=args.batch_size,
-        max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
-        enforce_eager=False,
-        disable_log_stats=True,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        # SSM state dtype (applies to both standard and ReplaySSM).
-        mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
-    )
-    if args.disable_flashinfer_autotune:
-        # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
-        # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
-        llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
-    if mode == "replayssm":
-        llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
-
-    _ssm_cm = None
-    if mode == "standard" and args.baseline_ssm_config:
-        from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
-
-        _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
-        _ssm_cm = override_ssm_config((_bsm, _nw))
-        _ssm_cm.__enter__()  # active through LLM() graph capture + decode
-        print(
-            f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
-            flush=True,
-        )
-
-    llm = LLM(**llm_kwargs)
-    prompts = [args.prompt] * args.batch_size
-
-    def timed_generate(n_tokens):
-        sp = SamplingParams(
-            n=1,
-            temperature=0.0,
-            ignore_eos=True,
-            min_tokens=n_tokens,
-            max_tokens=n_tokens,
-        )
-        if torch.accelerator.is_available():
-            torch.accelerator.synchronize()
-        t0 = time.perf_counter()
-        outs = llm.generate(prompts, sp, use_tqdm=False)
-        if torch.accelerator.is_available():
-            torch.accelerator.synchronize()
-        elapsed = time.perf_counter() - t0
-        produced = min(len(o.outputs[0].token_ids) for o in outs)
-        assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
-        return elapsed
-
-    timed_generate(args.warmup_steps)
-
-    best = None
-    for _ in range(args.repeats):
-        elapsed = timed_generate(args.num_steps)
-        tok_s = args.batch_size * args.num_steps / elapsed
-        per_step_ms = elapsed / args.num_steps * 1e3
-        print(
-            f"[{mode}] {elapsed:.3f}s  {tok_s:,.0f} tok/s  {per_step_ms:.3f} ms/step",
-            flush=True,
-        )
-        if best is None or elapsed < best["elapsed_s"]:
-            best = {
-                "mode": mode,
-                "elapsed_s": elapsed,
-                "tok_s": tok_s,
-                "per_step_ms": per_step_ms,
-            }
-
-    print("RESULT_JSON " + json.dumps(best), flush=True)
-    if _ssm_cm is not None:
-        _ssm_cm.__exit__(None, None, None)
-
-
-def run_one_mode(args, mode) -> dict:
-    cmd = [
-        sys.executable,
-        __file__,
-        "--worker",
-        mode,
-        "--model-id",
-        args.model_id,
-        "--prompt",
-        args.prompt,
-        "--batch-size",
-        str(args.batch_size),
-        "--num-steps",
-        str(args.num_steps),
-        "--warmup-steps",
-        str(args.warmup_steps),
-        "--repeats",
-        str(args.repeats),
-        "--buffer-len",
-        str(args.buffer_len),
-        "--dtype",
-        args.dtype,
-        "--gpu-memory-utilization",
-        str(args.gpu_memory_utilization),
-        "--mamba-ssm-cache-dtype",
-        args.mamba_ssm_cache_dtype,
-        "--baseline-ssm-config",
-        args.baseline_ssm_config,
-    ]
-    cmd.append(
-        "--disable-flashinfer-autotune"
-        if args.disable_flashinfer_autotune
-        else "--no-disable-flashinfer-autotune"
-    )
-    if args.max_model_len is not None:
-        cmd += ["--max-model-len", str(args.max_model_len)]
-
-    result = None
-    proc = subprocess.Popen(
-        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
-    )
-    for line in proc.stdout:
-        sys.stdout.write(line)
-        sys.stdout.flush()
-        if line.startswith("RESULT_JSON "):
-            result = json.loads(line[len("RESULT_JSON ") :])
-    proc.wait()
-    if proc.returncode != 0:
-        raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
-    if result is None:
-        raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
-    return result
-
-
-def main():
-    args = parse_args()
-    if args.worker is not None:
-        run_worker(args)
-        return
-
-    print(
-        f"model={args.model_id}  batch_size={args.batch_size}  "
-        f"steps={args.num_steps}  buffer_len={args.buffer_len}  dtype={args.dtype}"
-    )
-
-    std = run_one_mode(args, "standard")
-    fla = run_one_mode(args, "replayssm")
-    speedup = std["per_step_ms"] / fla["per_step_ms"]
-
-    print()
-    header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
-    print(header)
-    print("-" * len(header))
-    for r in (std, fla):
-        print(
-            f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
-            f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
-        )
-    print("-" * len(header))
-    print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
-
-
-if __name__ == "__main__":
-    main()
@@ -15,7 +15,6 @@ endif()
 #
 set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
-set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM})
 set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})

 include_directories("${CMAKE_SOURCE_DIR}/csrc")
@@ -97,14 +96,12 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
    set(ENABLE_NUMA OFF)
    check_sysctl(hw.optional.neon ASIMD_FOUND)
    check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
-    check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND)
 else()
    find_isa(${CPUINFO} "Power11" POWER11_FOUND)
    find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
    find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
    find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
    find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
-    find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support
    find_isa(${CPUINFO} "S390" S390_FOUND)
    find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
    find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
@@ -114,11 +111,6 @@ else()
        set(ARM_BF16_FOUND ON)
        message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
    endif()
-    if (ENABLE_ARM_I8MM)
-        set(ARM_I8MM_FOUND ON)
-        message(STATUS
-            "ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable")
-    endif()
    # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
    # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
    # overrides the detection result.
@@ -174,11 +166,6 @@ elseif (ASIMD_FOUND)
        message(WARNING "BF16 functionality is not available")
        set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")  
    endif()
-    if(ARM_I8MM_FOUND)
-        message(STATUS "I8MM extension detected")
-        string(APPEND MARCH_FLAGS "+i8mm")
-        add_compile_definitions(ARM_I8MM_SUPPORT)
-    endif()
    list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})     
 elseif (S390_FOUND)
    message(STATUS "S390 detected")
@@ -443,7 +430,6 @@ set(VLLM_EXT_SRC
    "csrc/cpu/layernorm.cpp"
    "csrc/cpu/mla_decode.cpp"
    "csrc/cpu/pos_encoding.cpp"
-    "csrc/cpu/mamba_cpu.cpp"
    "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")
@@ -460,13 +446,8 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
        "csrc/cpu/shm.cpp"
        "csrc/cpu/activation_lut_bf16.cpp"
        "csrc/cpu/cpu_tanhf_neon.hpp"
+        "csrc/cpu/cpu_fused_moe.cpp"
        ${VLLM_EXT_SRC})
-    if (ARM_BF16_FOUND)
-        set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
-        if (ARM_I8MM_FOUND)
-            set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
-        endif()
-    endif()
 endif()

 if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)	
@@ -508,7 +489,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
-        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -522,7 +502,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
@@ -68,9 +68,6 @@ endif()
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-      list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
-    endif()
  else()
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
  endif()
@@ -1,74 +0,0 @@
-include(FetchContent)
-
-if(DEFINED ENV{FLASH_KDA_SRC_DIR})
-  set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
-endif()
-
-if(FLASH_KDA_SRC_DIR)
-  FetchContent_Declare(
-    flashkda
-    SOURCE_DIR ${FLASH_KDA_SRC_DIR}
-  )
-else()
-  FetchContent_Declare(
-    flashkda
-    GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
-    GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
-    GIT_PROGRESS TRUE
-    GIT_SUBMODULES cutlass
-  )
-endif()
-
-FetchContent_MakeAvailable(flashkda)
-message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
-
-set(FLASH_KDA_SUPPORT_ARCHS)
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
-endif()
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
-elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
-endif()
-
-cuda_archs_loose_intersection(
-  FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
-
-if(FLASH_KDA_ARCHS)
-  message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
-
-  set(FLASH_KDA_SOURCES
-    csrc/flashkda_registration.cpp
-    ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
-    ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
-  set(FLASH_KDA_INCLUDES
-    ${flashkda_SOURCE_DIR}/csrc
-    ${flashkda_SOURCE_DIR}/cutlass/include
-    ${flashkda_SOURCE_DIR}/cutlass/examples/common
-    ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
-
-  set_gencode_flags_for_srcs(
-    SRCS "${FLASH_KDA_SOURCES}"
-    CUDA_ARCHS "${FLASH_KDA_ARCHS}")
-
-  define_extension_target(
-    _flashkda_C
-    DESTINATION vllm
-    LANGUAGE ${VLLM_GPU_LANG}
-    SOURCES ${FLASH_KDA_SOURCES}
-    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-    ARCHITECTURES ${VLLM_GPU_ARCHES}
-    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
-    USE_SABI 3
-    WITH_SOABI)
-
-  target_compile_options(_flashkda_C PRIVATE
-    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
-    $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
-else()
-  message(STATUS
-    "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
-    "(SM90, SM10x, or SM12x) are required")
-  add_custom_target(_flashkda_C)
-endif()
@@ -60,9 +60,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    # CUDA 12.9 has introduced "Family-Specific Architecture Features"
    # this supports all compute_10x family
    list(APPEND SUPPORT_ARCHS "10.0f")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-        list(APPEND SUPPORT_ARCHS "10.7f")
-    endif()
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
    list(APPEND SUPPORT_ARCHS "10.0a")
 endif()
@@ -191,3 +188,4 @@ else()
    add_custom_target(_flashmla_C)
    add_custom_target(_flashmla_extension_C)
 endif()
+
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
+    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
-  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")

  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _qutlass_fc_root)
@@ -55,11 +55,7 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")

 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
 else()
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
@@ -129,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )

+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
  define_extension_target(
    _qutlass_C
    DESTINATION vllm
@@ -141,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    WITH_SOABI)

  target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_MINIMAL_BUILD=1
+    QUTLASS_DISABLE_PYBIND=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
-    TORCH_TARGET_VERSION=0x020B000000000000ULL
-    USE_CUDA)
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -14,7 +14,7 @@ else()
  FetchContent_Declare(
    tml_fa4
    GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
-    GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
+    GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND "")
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
+          GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -241,15 +241,14 @@ endmacro()
 # `<major>.<minor>`, dedupes them and then sorts them in ascending order and 
 # stores them in `OUT_ARCHES`.
 #
-# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags.
-# This handles mismatches such as `arch=compute_20,code=sm_121`.
+# Example:
+#   CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" 
+#   extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS)
+#   OUT_ARCHES="7.5;...;9.0"
 function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
  set(_CUDA_ARCHES)
  foreach(_ARCH ${CUDA_ARCH_FLAGS})
-    string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
-    if (NOT _COMPUTE)
-      string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
-    endif()
+    string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
    if (_COMPUTE)
      set(_COMPUTE ${CMAKE_MATCH_1})
    endif()
@@ -397,24 +396,14 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
  # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
  # family. The output uses TGT's value to preserve the user's compilation flags.
  set(_CUDA_ARCHS)
-  # Resolve exact base matches before family fallbacks so a generic entry such
-  # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
-  foreach(_arch ${_SRC_CUDA_ARCHS})
-    if(_arch MATCHES "[af]$")
-      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
-        list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
-        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
-        list(APPEND _CUDA_ARCHS "${_arch}")
-      endif()
-    endif()
-  endforeach()
-
  foreach(_arch ${_SRC_CUDA_ARCHS})
    if(_arch MATCHES "[af]$")
      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
+      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
+        list(APPEND _CUDA_ARCHS "${_arch}")
+      elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
        list(APPEND _CUDA_ARCHS "${_base}a")
      elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -498,7 +487,7 @@ endfunction()

 function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
  endif()
@@ -67,8 +67,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
    torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
    torch::Tensor const& dst,               // [TOT_TOKENS, 576]
    torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& seq_lens,          // [BATCH]
    torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
+    int64_t batch_size);

 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(
@@ -172,15 +172,4 @@

 #endif  // __riscv_v

-// Power VSX
-#ifdef __powerpc__
-  // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
-  // implements a vectorised 5-term minimax polynomial using VSX intrinsics.
-  #define DEFINE_FAST_EXP                                                     \
-    auto fast_exp = [&](const vec_op::FP32Vec16& vec)                         \
-                        __attribute__((always_inline)) { return vec.exp(); }; \
-    auto fast_exp_f16 = fast_exp;
-
-#endif  // __powerpc__
-
 #endif
@@ -30,8 +30,7 @@ torch::Tensor get_scheduler_metadata(
    const torch::Tensor& query_start_loc, const bool causal,
    const int64_t window_size, const std::string& isa_hint,
    const bool enable_kv_split,
-    const std::optional<torch::Tensor>& dynamic_causal,
-    const std::string& kv_cache_dtype) {
+    const std::optional<torch::Tensor>& dynamic_causal) {
  cpu_attention::ISA isa;
  if (isa_hint == "amx") {
    isa = cpu_attention::ISA::AMX;
@@ -66,11 +65,9 @@ torch::Tensor get_scheduler_metadata(
  input.dynamic_causal =
      dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;

-  const int64_t kv_cache_idx =
-      static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
  VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
-    CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
-      input.elem_size = sizeof(attn_impl::kv_cache_t);
+    CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
+      input.elem_size = sizeof(scalar_t);
      input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
      input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
      input.output_buffer_elem_size =
@@ -102,9 +102,7 @@ class TileGemm82 {
    kv_cache_t* __restrict__ curr_b = b_tile;

    for (int32_t k = 0; k < dynamic_k_size; ++k) {
-      auto fp32_b_regs = load_b_pair_vec(curr_b);
-      auto fp32_b_0_reg = fp32_b_regs.first;
-      auto fp32_b_1_reg = fp32_b_regs.second;
+      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);

      float* __restrict__ curr_m_a = curr_a;
      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -1,6 +1,5 @@
 #include "cpu/cpu_types.hpp"
 #include "cpu/utils.hpp"
-#include "cpu/cpu_fused_moe_activations.hpp"
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
 #include "cpu/cpu_arch_macros.h"

@@ -44,9 +43,193 @@
  }()

 namespace {
+enum class FusedMOEAct {
+  SiluAndMul,
+  SwigluOAIAndMul,
+  GeluAndMul,
+  GeluTanhAndMul,
+};

-using cpu_fused_moe_utils::apply_gated_act;
-using cpu_fused_moe_utils::FusedMOEAct;
+FusedMOEAct get_act_type(const std::string& act) {
+  if (act == "silu") {
+    return FusedMOEAct::SiluAndMul;
+  } else if (act == "swigluoai") {
+    return FusedMOEAct::SwigluOAIAndMul;
+  } else if (act == "gelu") {
+    return FusedMOEAct::GeluAndMul;
+  } else if (act == "gelu_tanh") {
+    return FusedMOEAct::GeluTanhAndMul;
+  } else {
+    TORCH_CHECK(false, "Invalid act type: " + act);
+  }
+}
+
+template <typename scalar_t>
+void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
+                       const int32_t m_size, const int32_t n_size,
+                       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);
+  vec_op::FP32Vec16 one_vec(1.0);
+
+  DEFINE_FAST_EXP
+
+  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));
+      auto glu = gate_vec * sigmoid_vec;
+      auto gated_output_fp32 = (one_vec + up_vec) * glu;
+      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
+      gated_output.save(output + n / 2);
+    }
+    input += input_stride;
+    output += output_stride;
+  }
+}
+
+template <typename scalar_t>
+void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
+                  const int32_t m_size, const int32_t n_size,
+                  const int32_t input_stride, const int32_t output_stride) {
+  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
+  const int32_t dim = n_size / 2;
+  float* __restrict__ gate = input;
+  float* __restrict__ up = input + dim;
+  vec_op::FP32Vec16 one_vec(1.0);
+
+  DEFINE_FAST_EXP
+
+  for (int32_t m = 0; m < m_size; ++m) {
+    for (int32_t n = 0; n < dim; n += 16) {
+      vec_op::FP32Vec16 gate_vec(gate + n);
+      vec_op::FP32Vec16 up_vec(up + n);
+      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
+      auto silu = gate_vec * sigmoid_vec;
+      auto gated_output_fp32 = up_vec * silu;
+      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
+      gated_output.save(output + n);
+    }
+    gate += input_stride;
+    up += input_stride;
+    output += output_stride;
+  }
+}
+
+template <typename scalar_t>
+void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
+                  const int32_t m_size, const int32_t n_size,
+                  const int32_t input_stride, const int32_t output_stride) {
+  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
+  const int32_t dim = n_size / 2;
+  float* __restrict__ gate = input;
+  float* __restrict__ up = input + dim;
+  vec_op::FP32Vec16 one_vec(1.0);
+  vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
+  vec_op::FP32Vec16 w2_vec(0.5);
+  alignas(64) float temp[16];
+
+  DEFINE_FAST_EXP
+
+  for (int32_t m = 0; m < m_size; ++m) {
+    for (int32_t n = 0; n < dim; n += 16) {
+      vec_op::FP32Vec16 gate_vec(gate + n);
+      vec_op::FP32Vec16 up_vec(up + n);
+      auto er_input_vec = gate_vec * w1_vec;
+
+      er_input_vec.save(temp);
+      for (int32_t i = 0; i < 16; ++i) {
+        temp[i] = std::erf(temp[i]);
+      }
+      vec_op::FP32Vec16 er_vec(temp);
+      auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
+      auto gated_output_fp32 = up_vec * gelu;
+      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
+      gated_output.save(output + n);
+    }
+    gate += input_stride;
+    up += input_stride;
+    output += output_stride;
+  }
+}
+
+template <typename scalar_t>
+void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
+                       const int32_t m_size, const int32_t n_size,
+                       const int32_t input_stride,
+                       const int32_t output_stride) {
+  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
+  const int32_t dim = n_size / 2;
+  float* __restrict__ gate = input;
+  float* __restrict__ up = input + dim;
+  vec_op::FP32Vec16 one_vec(1.0);
+  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
+  vec_op::FP32Vec16 w2_vec(0.5);
+  vec_op::FP32Vec16 w3_vec(0.044715);
+
+  for (int32_t m = 0; m < m_size; ++m) {
+    for (int32_t n = 0; n < dim; n += 16) {
+      vec_op::FP32Vec16 gate_vec(gate + n);
+      vec_op::FP32Vec16 up_vec(up + n);
+      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
+      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();
+      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);
+      gated_output.save(output + n);
+    }
+    gate += input_stride;
+    up += input_stride;
+    output += output_stride;
+  }
+}
+
+template <typename scalar_t>
+FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
+                                  float* __restrict__ input,
+                                  scalar_t* __restrict__ output,
+                                  const int32_t m, const int32_t n,
+                                  const int32_t input_stride,
+                                  const int32_t output_stride) {
+  switch (act) {
+    case FusedMOEAct::SwigluOAIAndMul:
+      swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
+      return;
+    case FusedMOEAct::SiluAndMul:
+      silu_and_mul(input, output, m, n, input_stride, output_stride);
+      return;
+    case FusedMOEAct::GeluAndMul:
+      gelu_and_mul(input, output, m, n, input_stride, output_stride);
+      return;
+    case FusedMOEAct::GeluTanhAndMul:
+      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
+      return;
+    default:
+      TORCH_CHECK(false, "Unsupported act type.");
+  }
+}

 template <typename scalar_t, typename gemm_t>
 void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -634,7 +817,6 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
    }
  }
 }
-
 }  // namespace

 void prepack_moe_weight(
@@ -682,7 +864,7 @@ void cpu_fused_moe(
  const int32_t input_size_2 = w2.size(2);
  const int32_t output_size_2 = w2.size(1);
  const int32_t topk_num = topk_id.size(1);
-  const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
+  const FusedMOEAct act_type = get_act_type(act);
  cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
  TORCH_CHECK(!skip_weighted || topk_num == 1,
              "skip_weighted is only supported for topk=1 on CPU");
@@ -1,204 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
-#define CPU_FUSED_MOE_ACTIVATIONS_HPP
-
-#include <cmath>
-#include <cstdint>
-#include <string>
-
-#include "cpu/cpu_arch_macros.h"
-#include "cpu/utils.hpp"
-
-namespace cpu_fused_moe_utils {
-enum class FusedMOEAct {
-  SiluAndMul,
-  SwigluOAIAndMul,
-  GeluAndMul,
-  GeluTanhAndMul,
-};
-
-inline FusedMOEAct get_act_type(const std::string& act) {
-  if (act == "silu") {
-    return FusedMOEAct::SiluAndMul;
-  } else if (act == "swigluoai") {
-    return FusedMOEAct::SwigluOAIAndMul;
-  } else if (act == "gelu") {
-    return FusedMOEAct::GeluAndMul;
-  } else if (act == "gelu_tanh") {
-    return FusedMOEAct::GeluTanhAndMul;
-  } else {
-    TORCH_CHECK(false, "Invalid act type: " + act);
-  }
-}
-
-template <typename scalar_t>
-void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                       const int32_t m_size, const int32_t n_size,
-                       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);
-  vec_op::FP32Vec16 one_vec(1.0);
-
-  DEFINE_FAST_EXP
-
-  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));
-      auto glu = gate_vec * sigmoid_vec;
-      auto gated_output_fp32 = (one_vec + up_vec) * glu;
-      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
-      gated_output.save(output + n / 2);
-    }
-    input += input_stride;
-    output += output_stride;
-  }
-}
-
-template <typename scalar_t>
-void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                  const int32_t m_size, const int32_t n_size,
-                  const int32_t input_stride, const int32_t output_stride) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  const int32_t dim = n_size / 2;
-  float* __restrict__ gate = input;
-  float* __restrict__ up = input + dim;
-  vec_op::FP32Vec16 one_vec(1.0);
-
-  DEFINE_FAST_EXP
-
-  for (int32_t m = 0; m < m_size; ++m) {
-    for (int32_t n = 0; n < dim; n += 16) {
-      vec_op::FP32Vec16 gate_vec(gate + n);
-      vec_op::FP32Vec16 up_vec(up + n);
-      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
-      auto silu = gate_vec * sigmoid_vec;
-      auto gated_output_fp32 = up_vec * silu;
-      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
-      gated_output.save(output + n);
-    }
-    gate += input_stride;
-    up += input_stride;
-    output += output_stride;
-  }
-}
-
-template <typename scalar_t>
-void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                  const int32_t m_size, const int32_t n_size,
-                  const int32_t input_stride, const int32_t output_stride) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  const int32_t dim = n_size / 2;
-  float* __restrict__ gate = input;
-  float* __restrict__ up = input + dim;
-  vec_op::FP32Vec16 one_vec(1.0);
-  vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
-  vec_op::FP32Vec16 w2_vec(0.5);
-  alignas(64) float temp[16];
-
-  DEFINE_FAST_EXP
-
-  for (int32_t m = 0; m < m_size; ++m) {
-    for (int32_t n = 0; n < dim; n += 16) {
-      vec_op::FP32Vec16 gate_vec(gate + n);
-      vec_op::FP32Vec16 up_vec(up + n);
-      auto er_input_vec = gate_vec * w1_vec;
-
-      er_input_vec.save(temp);
-      for (int32_t i = 0; i < 16; ++i) {
-        temp[i] = std::erf(temp[i]);
-      }
-      vec_op::FP32Vec16 er_vec(temp);
-      auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
-      auto gated_output_fp32 = up_vec * gelu;
-      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
-      gated_output.save(output + n);
-    }
-    gate += input_stride;
-    up += input_stride;
-    output += output_stride;
-  }
-}
-
-template <typename scalar_t>
-void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                       const int32_t m_size, const int32_t n_size,
-                       const int32_t input_stride,
-                       const int32_t output_stride) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  const int32_t dim = n_size / 2;
-  float* __restrict__ gate = input;
-  float* __restrict__ up = input + dim;
-  vec_op::FP32Vec16 one_vec(1.0);
-  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
-  vec_op::FP32Vec16 w2_vec(0.5);
-  vec_op::FP32Vec16 w3_vec(0.044715);
-
-  for (int32_t m = 0; m < m_size; ++m) {
-    for (int32_t n = 0; n < dim; n += 16) {
-      vec_op::FP32Vec16 gate_vec(gate + n);
-      vec_op::FP32Vec16 up_vec(up + n);
-      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
-      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();
-      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);
-      gated_output.save(output + n);
-    }
-    gate += input_stride;
-    up += input_stride;
-    output += output_stride;
-  }
-}
-
-template <typename scalar_t>
-FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
-                                  float* __restrict__ input,
-                                  scalar_t* __restrict__ output,
-                                  const int32_t m, const int32_t n,
-                                  const int32_t input_stride,
-                                  const int32_t output_stride) {
-  switch (act) {
-    case FusedMOEAct::SwigluOAIAndMul:
-      swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
-    case FusedMOEAct::SiluAndMul:
-      silu_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
-    case FusedMOEAct::GeluAndMul:
-      gelu_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
-    case FusedMOEAct::GeluTanhAndMul:
-      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
-    default:
-      TORCH_CHECK(false, "Unsupported act type.");
-  }
-}
-}  // namespace cpu_fused_moe_utils
-
-#endif
@@ -1,647 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "cpu/cpu_arch_macros.h"
-
-#include <algorithm>
-#include <cstdint>
-#include <cstring>
-#include <optional>
-#include <string>
-
-#include "cpu/cpu_fused_moe_activations.hpp"
-#include "cpu/cpu_types.hpp"
-#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
-#include "cpu/utils.hpp"
-
-#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT)
-  #include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp"
-  #define NEON_DISPATCH(SCALAR_TYPE, ...)                                   \
-    case cpu_utils::ISA::NEON: {                                            \
-      using gemm_t =                                                        \
-          cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::NEON, SCALAR_TYPE>; \
-      return __VA_ARGS__();                                                 \
-    }
-#else
-  #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
-#endif
-
-#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
-  [&] {                                                        \
-    switch (ISA_TYPE) {                                        \
-      NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__)                  \
-      default: {                                               \
-        TORCH_CHECK(false, "Invalid CPU ISA type.");           \
-      }                                                        \
-    }                                                          \
-  }()
-
-namespace {
-using cpu_fused_moe_utils::apply_gated_act;
-using cpu_fused_moe_utils::FusedMOEAct;
-
-template <typename gemm_t>
-void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr,
-                                  int8_t* __restrict__ packed_weight_ptr,
-                                  const int32_t expert_num,
-                                  const int32_t output_size,
-                                  const int32_t input_size,
-                                  const int64_t expert_stride) {
-#pragma omp parallel for
-  for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
-    gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
-                        packed_weight_ptr + expert_stride * e_idx, output_size,
-                        input_size);
-  }
-}
-
-// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp
-template <typename scalar_t, typename gemm_t>
-void fused_moe_int8_impl(
-    scalar_t* __restrict__ output, const scalar_t* __restrict__ input,
-    const int8_t* __restrict__ w13, const int8_t* __restrict__ w2,
-    const float* __restrict__ w13_scales, const float* __restrict__ w2_scales,
-    scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias,
-    const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id,
-    const FusedMOEAct act_type, const int32_t token_num,
-    const int32_t expert_num, const int32_t topk_num,
-    const int32_t input_size_13, const int32_t output_size_13,
-    const int32_t input_size_2, const int32_t output_size_2,
-    const bool skip_weighted) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  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;
-
-  TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0);
-  TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 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 w13_input_buffer_size = cpu_utils::round_up<64>(
-      gemm_m_tile_size * input_size_13 * sizeof(int8_t));
-  const int32_t w2_input_buffer_size =
-      cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
-
-  const int32_t w13_n_tile_size = [&]() {
-    const int64_t cache_size = cpu_utils::get_available_l2_size();
-    const int32_t n_size_cache_limit =
-        (cache_size - w13_input_buffer_size) /
-        (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t));
-    const int32_t n_size_thread_limit =
-        output_size_13 / std::max(1, thread_num / topk_num);
-    const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
-        std::min(n_size_cache_limit, n_size_thread_limit));
-    return std::max(n_size, min_w13_n_tile_size);
-  }();
-
-  const int32_t w2_n_tile_size = [&]() {
-    const int64_t cache_size = cpu_utils::get_available_l2_size();
-    const int32_t n_size_cache_limit =
-        (cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_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>(
-        std::min(n_size_cache_limit, n_size_thread_limit));
-    return std::max(n_size, gemm_n_tile_size);
-  }();
-
-  int32_t common_buffer_offset = 0;
-  const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
-  common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
-  const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
-  common_buffer_offset +=
-      cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
-  const int32_t expanded_token_num = token_num * topk_num;
-  const int32_t expand_token_id_buffer_offset = common_buffer_offset;
-  common_buffer_offset +=
-      cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
-  const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
-  common_buffer_offset +=
-      cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
-  const int32_t input_quant_buffer_offset = common_buffer_offset;
-  common_buffer_offset +=
-      cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t));
-  const int32_t input_scale_buffer_offset = common_buffer_offset;
-  common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float));
-  const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
-  common_buffer_offset += cpu_utils::round_up<64>(
-      expanded_token_num * input_size_2 * sizeof(scalar_t));
-  const int32_t w13_output_scale_buffer_offset = common_buffer_offset;
-  common_buffer_offset +=
-      cpu_utils::round_up<64>(expanded_token_num * sizeof(float));
-  const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
-  common_buffer_offset += cpu_utils::round_up<64>(
-      expanded_token_num * output_size_2 * sizeof(float));
-
-  int32_t gemm_thread_buffer_offset = 0;
-  const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset;
-  gemm_thread_buffer_offset +=
-      std::max(w13_input_buffer_size, w2_input_buffer_size);
-  const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset;
-  gemm_thread_buffer_offset += cpu_utils::round_up<64>(
-      gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) *
-      sizeof(int32_t));
-
-  const int32_t ws_output_buffer_offset = 0;
-  const int32_t ws_thread_buffer_size =
-      cpu_utils::round_up<64>(output_size_2 * sizeof(float));
-  const int32_t thread_buffer_size =
-      std::max(gemm_thread_buffer_offset, ws_thread_buffer_size);
-  const int32_t buffer_size =
-      common_buffer_offset + thread_buffer_size * thread_num;
-  cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
-  uint8_t* common_buffer_start =
-      cpu_utils::ScratchPadManager::get_scratchpad_manager()
-          ->get_data<uint8_t>();
-  uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
-
-  int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
-      common_buffer_start + token_num_per_group_buffer_offset);
-  int32_t* __restrict__ cu_token_num_per_group_buffer =
-      reinterpret_cast<int32_t*>(common_buffer_start +
-                                 cu_token_num_per_group_buffer_offset);
-  int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
-      common_buffer_start + expand_token_id_buffer_offset);
-  int32_t* __restrict__ expand_token_id_index_buffer =
-      reinterpret_cast<int32_t*>(common_buffer_start +
-                                 expand_token_id_index_buffer_offset);
-  int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>(
-      common_buffer_start + input_quant_buffer_offset);
-  float* __restrict__ input_scale_buffer =
-      reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset);
-
-  std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
-  for (int32_t i = 0; i < expanded_token_num; ++i) {
-    ++token_num_per_group_buffer[topk_id[i]];
-  }
-
-  int32_t token_num_sum = 0;
-  cu_token_num_per_group_buffer[0] = 0;
-  int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
-  for (int32_t i = 0; i < expert_num; ++i) {
-    token_index_buffer[i] = token_num_sum;
-    token_num_sum += token_num_per_group_buffer[i];
-  }
-
-  for (int32_t i = 0; i < token_num; ++i) {
-    const int32_t* curr_topk_id = topk_id + i * topk_num;
-    int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
-    for (int32_t j = 0; j < topk_num; ++j) {
-      const int32_t curr_expert_id = curr_topk_id[j];
-      const int32_t curr_index = token_index_buffer[curr_expert_id]++;
-      expand_token_id_buffer[curr_index] = i;
-      curr_index_buffer[j] = curr_index;
-    }
-  }
-
-// quantize inputs
-#pragma omp parallel for
-  for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
-    gemm_t::quantize_row(input + token_idx * input_size_13,
-                         input_quant_buffer + token_idx * input_size_13,
-                         input_scale_buffer[token_idx], input_size_13);
-  }
-
-  {
-    alignas(64) cpu_utils::Counter counter;
-    cpu_utils::Counter* counter_ptr = &counter;
-
-// w13 GEMM + act
-#pragma omp parallel for schedule(static, 1)
-    for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
-      const int32_t task_num_per_expert =
-          (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
-      const int32_t task_num = task_num_per_expert * expert_num;
-      uint8_t* __restrict__ thread_buffer =
-          thread_buffer_start + thread_id * thread_buffer_size;
-      int8_t* __restrict__ gemm_input_buffer =
-          reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
-      float* __restrict__ gemm_output_buffer =
-          reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
-      auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
-          common_buffer_start + w13_gemm_output_buffer_offset);
-      gemm_t gemm;
-
-      const int32_t w13_n_group_stride =
-          gemm_t::WeightOCGroupSize * input_size_13;
-      const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
-
-      for (;;) {
-        const int32_t task_id = counter_ptr->acquire_counter();
-        if (task_id >= task_num) {
-          break;
-        }
-        const int32_t curr_expert_id = task_id / task_num_per_expert;
-        const int32_t curr_output_group_id = task_id % task_num_per_expert;
-        const int32_t curr_token_num =
-            token_num_per_group_buffer[curr_expert_id];
-        if (curr_token_num == 0) {
-          continue;
-        }
-
-        const int32_t actual_n_tile_size =
-            std::min(w13_n_tile_size,
-                     output_size_13 - curr_output_group_id * w13_n_tile_size);
-        const int32_t* __restrict__ curr_expand_token_id_buffer =
-            expand_token_id_buffer +
-            cu_token_num_per_group_buffer[curr_expert_id];
-        scalar_t* __restrict__ curr_w13_gemm_output_buffer =
-            w13_gemm_output_buffer +
-            cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 +
-            curr_output_group_id * w13_n_tile_size / 2;
-
-        const int8_t* w13_weight_ptr_0 = nullptr;
-        const int8_t* w13_weight_ptr_1 = nullptr;
-        const float* w13_scale_ptr_0 = nullptr;
-        const float* w13_scale_ptr_1 = nullptr;
-        scalar_t* w13_bias_ptr_0 = nullptr;
-        scalar_t* w13_bias_ptr_1 = nullptr;
-        if (act_type == FusedMOEAct::SwigluOAIAndMul) {
-          const int32_t output_offset = curr_output_group_id * w13_n_tile_size;
-          w13_weight_ptr_0 = w13 +
-                             curr_expert_id * input_size_13 * output_size_13 +
-                             output_offset * input_size_13;
-          w13_weight_ptr_1 =
-              w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
-          w13_scale_ptr_0 =
-              w13_scales + curr_expert_id * output_size_13 + output_offset;
-          w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2;
-          if (w13_bias != nullptr) {
-            w13_bias_ptr_0 =
-                w13_bias + curr_expert_id * output_size_13 + output_offset;
-            w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
-          }
-        } else {
-          const int32_t output_offset =
-              curr_output_group_id * (w13_n_tile_size / 2);
-          w13_weight_ptr_0 = w13 +
-                             curr_expert_id * input_size_13 * output_size_13 +
-                             output_offset * input_size_13;
-          w13_weight_ptr_1 =
-              w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
-          w13_scale_ptr_0 =
-              w13_scales + curr_expert_id * output_size_13 + output_offset;
-          w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2;
-          if (w13_bias != nullptr) {
-            w13_bias_ptr_0 =
-                w13_bias + curr_expert_id * output_size_13 + output_offset;
-            w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
-          }
-        }
-
-        for (int32_t token_idx = 0; token_idx < curr_token_num;
-             token_idx += gemm_m_tile_size) {
-          const int32_t actual_token_num =
-              std::min(gemm_m_tile_size, curr_token_num - token_idx);
-          const int8_t* input_rows[gemm_m_tile_size];
-          alignas(64) float input_scales[gemm_m_tile_size];
-          // gather and pack
-          for (int32_t i = 0; i < actual_token_num; ++i) {
-            const int32_t curr_token_id = curr_expand_token_id_buffer[i];
-            input_rows[i] = input_quant_buffer + curr_token_id * input_size_13;
-            input_scales[i] = input_scale_buffer[curr_token_id];
-          }
-          gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
-                                       actual_token_num, input_size_13);
-          curr_expand_token_id_buffer += actual_token_num;
-
-          const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0;
-          const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1;
-          const float* w13_scale_ptr_0_iter = w13_scale_ptr_0;
-          const float* w13_scale_ptr_1_iter = w13_scale_ptr_1;
-          scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0;
-          scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1;
-          float* w13_output_buffer_0_iter = gemm_output_buffer;
-          float* w13_output_buffer_1_iter =
-              gemm_output_buffer + actual_n_tile_size / 2;
-
-          for (int32_t i = 0; i < actual_n_tile_size;
-               i += min_w13_n_tile_size) {
-            auto* output_0_int32 =
-                reinterpret_cast<int32_t*>(w13_output_buffer_0_iter);
-            gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32,
-                      actual_token_num, input_size_13, w13_n_group_stride,
-                      actual_n_tile_size);
-            gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter,
-                                    input_scales, w13_scale_ptr_0_iter,
-                                    actual_token_num, gemm_n_tile_size,
-                                    actual_n_tile_size);
-            if (w13_bias != nullptr) {
-              cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
-                  w13_output_buffer_0_iter, w13_output_buffer_0_iter,
-                  w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
-                  actual_n_tile_size);
-              w13_bias_ptr_0_iter += gemm_n_tile_size;
-            }
-
-            auto* output_1_int32 =
-                reinterpret_cast<int32_t*>(w13_output_buffer_1_iter);
-            gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32,
-                      actual_token_num, input_size_13, w13_n_group_stride,
-                      actual_n_tile_size);
-            gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter,
-                                    input_scales, w13_scale_ptr_1_iter,
-                                    actual_token_num, gemm_n_tile_size,
-                                    actual_n_tile_size);
-            if (w13_bias != nullptr) {
-              cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
-                  w13_output_buffer_1_iter, w13_output_buffer_1_iter,
-                  w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
-                  actual_n_tile_size);
-              w13_bias_ptr_1_iter += gemm_n_tile_size;
-            }
-
-            w13_weight_ptr_0_iter += w13_n_tile_stride;
-            w13_weight_ptr_1_iter += w13_n_tile_stride;
-            w13_scale_ptr_0_iter += gemm_n_tile_size;
-            w13_scale_ptr_1_iter += gemm_n_tile_size;
-            w13_output_buffer_0_iter += gemm_n_tile_size;
-            w13_output_buffer_1_iter += gemm_n_tile_size;
-          }
-
-          apply_gated_act(act_type, gemm_output_buffer,
-                          curr_w13_gemm_output_buffer, actual_token_num,
-                          actual_n_tile_size, actual_n_tile_size, input_size_2);
-          curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
-        }
-      }
-    }
-  }
-
-  auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
-      common_buffer_start + w13_gemm_output_buffer_offset);
-  float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>(
-      common_buffer_start + w13_output_scale_buffer_offset);
-
-// quantize w2 inputs - in place
-#pragma omp parallel for
-  for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) {
-    scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2;
-    int8_t* output_row = reinterpret_cast<int8_t*>(input_row);
-    gemm_t::quantize_row(input_row, output_row,
-                         w13_output_scale_buffer[token_idx], input_size_2);
-  }
-
-  {
-    alignas(64) cpu_utils::Counter counter;
-    cpu_utils::Counter* counter_ptr = &counter;
-
-// w2 gemm
-#pragma omp parallel for schedule(static, 1)
-    for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
-      const int32_t task_num_per_expert =
-          (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
-      const int32_t task_num = task_num_per_expert * expert_num;
-      uint8_t* __restrict__ thread_buffer =
-          thread_buffer_start + thread_id * thread_buffer_size;
-      int8_t* __restrict__ gemm_input_buffer =
-          reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
-      float* __restrict__ gemm_output_buffer =
-          reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
-      float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
-          common_buffer_start + w2_gemm_output_buffer_offset);
-      gemm_t gemm;
-
-      const int32_t w2_n_group_stride =
-          gemm_t::WeightOCGroupSize * input_size_2;
-      const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
-
-      for (;;) {
-        const int32_t task_id = counter_ptr->acquire_counter();
-        if (task_id >= task_num) {
-          break;
-        }
-        const int32_t curr_expert_id = task_id / task_num_per_expert;
-        const int32_t curr_output_group_id = task_id % task_num_per_expert;
-        const int32_t curr_token_num =
-            token_num_per_group_buffer[curr_expert_id];
-        if (curr_token_num == 0) {
-          continue;
-        }
-
-        const int32_t actual_n_tile_size =
-            std::min(w2_n_tile_size,
-                     output_size_2 - curr_output_group_id * w2_n_tile_size);
-        scalar_t* __restrict__ curr_w13_gemm_output_buffer =
-            w13_gemm_output_buffer +
-            cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
-        float* __restrict__ curr_w13_output_scale_buffer =
-            w13_output_scale_buffer +
-            cu_token_num_per_group_buffer[curr_expert_id];
-        float* __restrict__ curr_w2_gemm_output_buffer =
-            w2_gemm_output_buffer +
-            cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
-            curr_output_group_id * w2_n_tile_size;
-        const int8_t* __restrict__ w2_weight_ptr =
-            w2 + curr_expert_id * output_size_2 * input_size_2 +
-            curr_output_group_id * w2_n_tile_size * input_size_2;
-        const float* __restrict__ w2_scale_ptr =
-            w2_scales + curr_expert_id * output_size_2 +
-            curr_output_group_id * w2_n_tile_size;
-        scalar_t* w2_bias_ptr = nullptr;
-        if (w2_bias != nullptr) {
-          w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
-                        curr_output_group_id * w2_n_tile_size;
-        }
-
-        for (int32_t token_idx = 0; token_idx < curr_token_num;
-             token_idx += gemm_m_tile_size) {
-          const int32_t actual_token_num =
-              std::min(gemm_m_tile_size, curr_token_num - token_idx);
-          const int8_t* input_rows[gemm_m_tile_size];
-          alignas(64) float input_scales[gemm_m_tile_size];
-          for (int32_t i = 0; i < actual_token_num; ++i) {
-            input_rows[i] = reinterpret_cast<const int8_t*>(
-                curr_w13_gemm_output_buffer + i * input_size_2);
-            input_scales[i] = curr_w13_output_scale_buffer[i];
-          }
-          gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
-                                       actual_token_num, input_size_2);
-
-          const int8_t* w2_weight_ptr_iter = w2_weight_ptr;
-          const float* w2_scale_ptr_iter = w2_scale_ptr;
-          scalar_t* w2_bias_ptr_iter = w2_bias_ptr;
-          float* 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) {
-            auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer);
-            gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32,
-                      actual_token_num, input_size_2, w2_n_group_stride,
-                      gemm_n_tile_size);
-            gemm_t::dequantize_tile(output_int32, gemm_output_buffer,
-                                    input_scales, w2_scale_ptr_iter,
-                                    actual_token_num, gemm_n_tile_size,
-                                    gemm_n_tile_size);
-            if (w2_bias != nullptr) {
-              cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
-                  gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter,
-                  actual_token_num, gemm_n_tile_size, gemm_n_tile_size);
-              w2_bias_ptr_iter += gemm_n_tile_size;
-            }
-            for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) {
-              std::memcpy(
-                  curr_w2_gemm_output_buffer_iter + m_idx * output_size_2,
-                  gemm_output_buffer + m_idx * gemm_n_tile_size,
-                  gemm_n_tile_size * sizeof(float));
-            }
-
-            w2_weight_ptr_iter += w2_n_tile_stride;
-            w2_scale_ptr_iter += gemm_n_tile_size;
-            curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
-          }
-
-          curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
-          curr_w13_output_scale_buffer += gemm_m_tile_size;
-          curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
-        }
-      }
-    }
-  }
-
-  {
-    alignas(64) cpu_utils::Counter counter;
-    cpu_utils::Counter* counter_ptr = &counter;
-
-#pragma omp parallel for schedule(static, 1)
-    for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
-      uint8_t* __restrict__ thread_buffer =
-          thread_buffer_start + thread_id * thread_buffer_size;
-      float* __restrict__ ws_output_buffer =
-          reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
-      float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
-          common_buffer_start + w2_gemm_output_buffer_offset);
-
-      for (;;) {
-        const int32_t token_id = counter_ptr->acquire_counter();
-        if (token_id >= token_num) {
-          break;
-        }
-        int32_t* __restrict__ curr_expand_token_id_index_buffer =
-            expand_token_id_index_buffer + token_id * topk_num;
-        const float* __restrict__ curr_weight =
-            topk_weights + token_id * topk_num;
-        const float first_weight = skip_weighted ? 1.0f : curr_weight[0];
-        scalar_t* __restrict__ curr_output_buffer =
-            output + token_id * output_size_2;
-
-        if (topk_num > 1) {
-          int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
-          float* w2_output_iter =
-              w2_gemm_output_buffer + w2_output_idx * output_size_2;
-          float* ws_output_buffer_iter = ws_output_buffer;
-          vec_op::FP32Vec16 weight_vec(first_weight);
-          for (int32_t i = 0; i < output_size_2; i += 16) {
-            vec_op::FP32Vec16 vec(w2_output_iter);
-            (vec * weight_vec).save(ws_output_buffer_iter);
-            w2_output_iter += 16;
-            ws_output_buffer_iter += 16;
-          }
-
-          for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
-            w2_output_idx = curr_expand_token_id_index_buffer[idx];
-            w2_output_iter =
-                w2_gemm_output_buffer + w2_output_idx * output_size_2;
-            ws_output_buffer_iter = ws_output_buffer;
-            weight_vec = vec_op::FP32Vec16(curr_weight[idx]);
-            for (int32_t i = 0; i < output_size_2; i += 16) {
-              vec_op::FP32Vec16 vec(w2_output_iter);
-              vec_op::FP32Vec16 sum(ws_output_buffer_iter);
-              (sum + vec * weight_vec).save(ws_output_buffer_iter);
-              w2_output_iter += 16;
-              ws_output_buffer_iter += 16;
-            }
-          }
-
-          const int32_t last_idx = topk_num - 1;
-          w2_output_idx = curr_expand_token_id_index_buffer[last_idx];
-          w2_output_iter =
-              w2_gemm_output_buffer + w2_output_idx * output_size_2;
-          ws_output_buffer_iter = ws_output_buffer;
-          scalar_t* curr_output_buffer_iter = curr_output_buffer;
-          weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]);
-          for (int32_t i = 0; i < output_size_2; i += 16) {
-            vec_op::FP32Vec16 vec(w2_output_iter);
-            vec_op::FP32Vec16 sum(ws_output_buffer_iter);
-            scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter);
-            w2_output_iter += 16;
-            ws_output_buffer_iter += 16;
-            curr_output_buffer_iter += 16;
-          }
-        } else {
-          const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
-          float* w2_output_iter =
-              w2_gemm_output_buffer + w2_output_idx * output_size_2;
-          scalar_t* curr_output_buffer_iter = curr_output_buffer;
-          vec_op::FP32Vec16 weight_vec(first_weight);
-          for (int32_t i = 0; i < output_size_2; i += 16) {
-            vec_op::FP32Vec16 vec(w2_output_iter);
-            scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter);
-            w2_output_iter += 16;
-            curr_output_buffer_iter += 16;
-          }
-        }
-      }
-    }
-  }
-}
-}  // namespace
-
-void prepack_moe_weight_int8(
-    const torch::Tensor& weight,  // [expert_num, output_size, input_size]
-    torch::Tensor& packed_weight, const std::string& isa) {
-  TORCH_CHECK(weight.is_contiguous());
-  const int32_t expert_num = weight.size(0);
-  const int32_t output_size = weight.size(1);
-  const int32_t input_size = weight.size(2);
-  const int64_t expert_stride = weight.stride(0);
-  const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
-  TORCH_CHECK_EQ(output_size % 32, 0);
-
-  CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() {
-    TORCH_CHECK_EQ(input_size % gemm_t::K, 0);
-    prepack_moe_weight_int8_impl<gemm_t>(
-        weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num,
-        output_size, input_size, expert_stride);
-  });
-}
-
-void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
-                        const torch::Tensor& w13, const torch::Tensor& w2,
-                        const torch::Tensor& w13_scale,
-                        const torch::Tensor& w2_scale,
-                        const std::optional<torch::Tensor>& w13_bias,
-                        const std::optional<torch::Tensor>& w2_bias,
-                        const torch::Tensor& topk_weights,
-                        const torch::Tensor& topk_id, const bool skip_weighted,
-                        const std::string& act, const std::string& isa) {
-  const int32_t token_num = input.size(0);
-  const int32_t input_size_13 = input.size(1);
-  const int64_t input_stride = input.stride(0);
-  TORCH_CHECK_EQ(input_stride, input_size_13);
-  const int32_t expert_num = w13.size(0);
-  const int32_t output_size_13 = w13.size(1);
-  const int32_t input_size_2 = w2.size(2);
-  const int32_t output_size_2 = w2.size(1);
-  const int32_t topk_num = topk_id.size(1);
-  const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
-  const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
-  TORCH_CHECK(!skip_weighted || topk_num == 1,
-              "skip_weighted is only supported for topk=1 on CPU");
-
-  VLLM_DISPATCH_FLOATING_TYPES(
-      input.scalar_type(), "cpu_fused_moe_int8", [&]() {
-        CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() {
-          fused_moe_int8_impl<scalar_t, gemm_t>(
-              output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-              w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(),
-              w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(),
-              w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
-              w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
-              topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(),
-              act_type, token_num, expert_num, topk_num, input_size_13,
-              output_size_13, input_size_2, output_size_2, skip_weighted);
-        });
-      });
-}
@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {

  explicit FP32Vec4(__vector float data) : reg(data) {}

-  FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
 };

 struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {

  explicit FP32Vec8(f32x4x2_t data) : reg(data) {}

-  FP32Vec8(const FP32Vec8& data) {
+  explicit FP32Vec8(const FP32Vec8& data) {
    reg.val[0] = data.reg.val[0];
    reg.val[1] = data.reg.val[1];
  }
@@ -336,14 +336,13 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
    reg.val[1] = fp16_to_fp32_bits(raw_lo);
  }
  float reduce_sum() const {
-    // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
-    // Step 1: pairwise sum of the two 4-wide halves
-    __vector float s = vec_add(reg.val[0], reg.val[1]);
-    // Step 2: rotate by 8 bytes (2 floats) and add
-    s = vec_add(s, vec_sld(s, s, 8));
-    // Step 3: rotate by 4 bytes (1 float) and add  => all lanes hold total
-    s = vec_add(s, vec_sld(s, s, 4));
-    return vec_extract(s, 0);
+    AliasReg ar;
+    ar.reg = reg;
+    float result = 0;
+    unroll_loop<int, VEC_ELEM_NUM>(
+        [&result, &ar](int i) { result += ar.values[i]; });
+
+    return result;
  }
  FP32Vec8 exp() const {
    f32x4x2_t out;
@@ -593,7 +592,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
  explicit FP32Vec16(f32x4x4_t data) : reg(data) {}

-  FP32Vec16(const FP32Vec16& data) {
+  explicit FP32Vec16(const FP32Vec16& data) {
    reg.val[0] = data.reg.val[0];
    reg.val[1] = data.reg.val[1];
    reg.val[2] = data.reg.val[2];
@@ -747,15 +746,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                                vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
  }

-  FP32Vec16 exp() const {
-    FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
-    FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
-    auto lo_e = lo.exp();
-    auto hi_e = hi.exp();
-    return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
-                               hi_e.reg.val[0], hi_e.reg.val[1]});
-  }
-
  float reduce_max() {
    __vector float max01 = vec_max(reg.val[0], reg.val[1]);
    __vector float max23 = vec_max(reg.val[2], reg.val[3]);
@@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {

  explicit FP32Vec4(__vector float data) : reg(data) {}

-  FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
 };

 struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -298,7 +298,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {

  explicit FP32Vec8(f32x4x2_t data) : reg(data) {}

-  FP32Vec8(const FP32Vec8& data) {
+  explicit FP32Vec8(const FP32Vec8& data) {
    reg.val[0] = data.reg.val[0];
    reg.val[1] = data.reg.val[1];
  }
@@ -643,7 +643,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  explicit FP32Vec16(f32x4x4_t data) : reg(data) {}

-  FP32Vec16(const FP32Vec16& data) {
+  explicit FP32Vec16(const FP32Vec16& data) {
    reg.val[0] = data.reg.val[0];
    reg.val[1] = data.reg.val[1];
    reg.val[2] = data.reg.val[2];
@@ -1,285 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
-// mamba_kernels.hpp.
-
-#include "cpu/mamba_kernels.hpp"
-
-#include <ATen/ATen.h>
-#include <torch/library.h>
-#include <c10/util/Optional.h>
-
-#include "cpu_types.hpp"
-
-// ---------------------------------------------------------------------------
-// causal_conv1d_update
-// ---------------------------------------------------------------------------
-at::Tensor causal_conv1d_update_cpu_impl(
-    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
-    const c10::optional<at::Tensor>& bias,
-    const c10::optional<std::string>& activation,
-    const c10::optional<at::Tensor>& conv_state_indices,
-    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
-  bool do_silu = false;
-  if (activation.has_value()) {
-    const std::string& act = activation.value();
-    do_silu = (act == "silu" || act == "swish");
-  }
-
-  at::ScalarType dtype = x.scalar_type();
-
-  // Input x: contiguous in native dtype.
-  at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
-
-  // conv_state: NEVER copy the full paged tensor just for layout reasons.
-  // If the dtype matches we work directly on conv_state (contiguous or not)
-  // by extracting strides and passing them to the kernel.
-  // Only a dtype-conversion copy is made when types differ (rare for BF16).
-  bool state_type_ok = (conv_state.scalar_type() == dtype);
-  at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
-  // state_c and conv_state may be non-contiguous — that is intentional.
-
-  // Weight: coerce to same dtype if needed (should match in practice)
-  at::Tensor w_c =
-      (weight.scalar_type() != dtype)
-          ? weight.to(dtype).contiguous()
-          : (weight.is_contiguous() ? weight : weight.contiguous());
-
-  // Bias stays float32 (small scalar, used only for fp32 accumulation)
-  at::Tensor bias_f32;
-  if (bias.has_value() && bias.value().defined())
-    bias_f32 = bias.value().to(at::kFloat).contiguous();
-
-  int64_t batch = x_c.size(0);
-  int64_t dim = x_c.size(1);
-  int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
-  int64_t width = w_c.size(1);
-  int64_t state_len = state_c.size(2);
-
-  // Extract strides — works for contiguous AND non-contiguous (transposed)
-  // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
-  // contiguous) stride(1): between conv channels (dim stride) stride(2):
-  // between state elements (=1 when contiguous, =dim when transposed)
-  int64_t stride_s_slot = state_c.stride(0);
-  int64_t stride_s_dim = state_c.stride(1);
-  int64_t stride_s_state = state_c.stride(2);
-
-  at::Tensor out = x_c.clone();  // native dtype, no float32 alloc
-
-  const int32_t* cache_idx_ptr = nullptr;
-  at::Tensor cache_idx_int;
-  if (conv_state_indices.has_value()) {
-    cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
-    cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
-  }
-
-  VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
-    mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
-        x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
-        stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
-        bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
-        out.data_ptr<scalar_t>(), cache_idx_ptr,
-        static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
-        do_silu);
-  });
-
-  // Write back only when a type-conversion copy was made.
-  // Layout-only non-contiguity is handled via strides above — no copy needed.
-  if (!state_type_ok) conv_state.copy_(state_c);
-
-  return out;
-}
-
-// ---------------------------------------------------------------------------
-// selective_state_update
-// ---------------------------------------------------------------------------
-void selective_state_update_cpu_impl(
-    at::Tensor& state,    // (nstates, nheads, dim, dstate)
-    const at::Tensor& x,  // (N, nheads, dim)
-    const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
-    const at::Tensor& C, const c10::optional<at::Tensor>& D,
-    const c10::optional<at::Tensor>& z,
-    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
-    const c10::optional<at::Tensor>& state_batch_indices,
-    const c10::optional<at::Tensor>& dst_state_batch_indices,
-    int64_t null_block_id, at::Tensor& out,
-    const c10::optional<at::Tensor>& num_accepted_tokens,
-    const c10::optional<at::Tensor>& cu_seqlens) {
-  at::ScalarType state_type = state.scalar_type();
-  at::ScalarType input_type = x.scalar_type();
-
-  // x, B, C must be contiguous and match input_type
-  auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor x_in = ensure_input(x);
-  at::Tensor B_in = ensure_input(B);
-  at::Tensor C_in = ensure_input(C);
-  at::Tensor z_in;
-  if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
-
-  // A, D, dt_bias are float32 model parameters that arrive here as expanded
-  // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
-  // We need just the scalar value per head as a (nheads,) 1-D array so that
-  // A_ptr[h] in the kernel correctly reads head h's value.
-  //
-  // Strategy: peel trailing expanded (stride=0) dims via .select(), which is
-  // a zero-copy view.  For A: (nheads, head_dim, dstate) strides (1,0,0)
-  //   → .select(2,0) → (nheads, head_dim) strides (1,0)
-  //   → .select(1,0) → (nheads,)          stride  (1,)   ← contiguous, free.
-  // No allocation, no type conversion (A is already float32).
-  auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = t;
-    // Peel trailing dimensions that are broadcast (stride=0 or size=1)
-    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
-    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-
-  at::Tensor A_f32 = to_per_head_1d_f32(A);  // (nheads,) float32
-  at::Tensor D_f32, dt_bias_f32;
-  if (D.has_value() && D.value().defined())
-    D_f32 = to_per_head_1d_f32(D.value());
-  if (dt_bias.has_value() && dt_bias.value().defined())
-    dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
-
-  // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
-  // the type conversion so we convert head_dim x fewer elements.
-  at::Tensor dt_f32;
-  {
-    // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
-    // take a zero-copy view of index 0 along that dim first.
-    at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt;  // (N, nheads)
-    at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
-    dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
-  }
-
-  int64_t nheads = state.size(1);
-  int64_t dim = state.size(2);
-  int64_t dstate = state.size(3);
-  int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
-                  ? cu_seqlens.value().size(0) - 1
-                  : x_in.size(0);
-  int64_t ngroups = B_in.size(1);
-
-  // Strides
-  int64_t stride_state_n = state.stride(0);
-  int64_t stride_state_h = state.stride(1);
-  int64_t stride_state_d = state.stride(2);
-  int64_t stride_x_n = x_in.stride(0);
-  int64_t stride_x_h = x_in.stride(1);
-  int64_t stride_dt_n = dt_f32.stride(0);  // dt is (N, nheads)
-  int64_t stride_BC_n = B_in.stride(0);
-  int64_t stride_BC_g = B_in.stride(1);
-  int64_t stride_out_n = out.stride(0);
-  int64_t stride_out_h = out.stride(1);
-
-  // Optional index pointers
-  auto get_int32_ptr =
-      [](const c10::optional<at::Tensor>& opt) -> const int32_t* {
-    return (opt.has_value() && opt.value().defined())
-               ? opt.value().data_ptr<int32_t>()
-               : nullptr;
-  };
-  const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
-  const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
-  const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
-  const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
-
-  // Dispatch on (state_t, input_t, out_t): write directly into `out`
-  // without any intermediate float32 buffer.
-  VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
-    using state_t = scalar_t;
-    VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
-      using input_t = scalar_t;
-      VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
-        using out_t = scalar_t;
-        mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
-            state.data_ptr<state_t>(), stride_state_n, stride_state_h,
-            stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
-            dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
-            B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
-            stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
-            z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
-            dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
-            out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
-            dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
-            nheads, ngroups, dim, dstate, dt_softplus);
-      });
-    });
-  });
-}
-
-// ---------------------------------------------------------------------------
-// mamba_chunk_scan_fwd_cpu
-// ---------------------------------------------------------------------------
-void mamba_chunk_scan_fwd_cpu_impl(
-    at::Tensor& out,  // [seqlen, nheads, headdim] — pre-allocated by caller
-    at::Tensor&
-        final_states,     // [batch, nheads, headdim, dstate] float32 contiguous
-    const at::Tensor& x,  // [seqlen, nheads, headdim]
-    const at::Tensor&
-        dt,  // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
-    const at::Tensor& A,                 // [nheads] float32
-    const at::Tensor& B,                 // [seqlen, ngroups, dstate]
-    const at::Tensor& C,                 // [seqlen, ngroups, dstate]
-    const c10::optional<at::Tensor>& D,  // [nheads] float32 (optional)
-    const c10::optional<at::Tensor>& z,  // [seqlen, nheads, headdim] (optional)
-    const at::Tensor& cu_seqlens         // [batch+1] int32
-) {
-  const at::ScalarType input_type = x.scalar_type();
-
-  auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor x_in = ensure_contig(x);
-  at::Tensor B_in = ensure_contig(B);
-  at::Tensor C_in = ensure_contig(C);
-  at::Tensor z_in;
-  if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
-
-  // A and D are float32 model parameters, potentially broadcast-expanded.
-  // Strip trailing broadcast dims to get a contiguous (nheads,) array.
-  auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = t;
-    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
-    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor A_f32 = to_per_head_f32(A);
-  at::Tensor D_f32;
-  if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
-
-  // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
-  // Python.
-  at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
-  if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
-
-  at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
-
-  const int64_t batch = final_states.size(0);
-  const int64_t nheads = final_states.size(1);
-  const int64_t headdim = final_states.size(2);
-  const int64_t dstate = final_states.size(3);
-  const int64_t ngroups = B_in.size(1);
-
-  TORCH_CHECK(final_states.is_contiguous(),
-              "mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
-  TORCH_CHECK(out.is_contiguous(),
-              "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
-              "raw data_ptr)");
-
-  VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
-    mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
-        final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
-        dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
-        B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
-        D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
-        z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
-        out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
-        ngroups, headdim, dstate);
-  });
-}
@@ -1,382 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// Fused CPU vector kernels for Mamba decode-step hotspots:
-//   - causal_conv1d_update  (depthwise 1-D conv state roll + compute)
-//   - selective_state_update (SSM recurrence, single-step)
-
-#pragma once
-
-#include "cpu_types.hpp"
-#include <cmath>
-#include <cstring>
-#include <cstdint>
-#include <algorithm>
-
-namespace mamba_cpu {
-
-// ---------------------------------------------------------------------------
-// causal_conv1d_update — templated for native BF16/FP32
-//
-// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
-// Explicit strides are passed so the kernel writes directly into the
-// correct memory locations without making a contiguous copy of the full
-// paged tensor (which was the source of the 34-41% direct_copy_kernel).
-//
-//   stride_s_slot  = state.stride(0)  — between cache slots
-//   stride_s_dim   = state.stride(1)  — between conv_dim channels
-//   stride_s_state = state.stride(2)  — between state elements
-//
-// When stride_s_state == 1 (contiguous), the memmove fast path is used.
-// ---------------------------------------------------------------------------
-template <typename scalar_t>
-inline void causal_conv1d_update_kernel(
-    const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
-    int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
-    const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
-    scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
-    int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
-    int64_t width, int64_t state_len, bool do_silu) {
-#pragma omp parallel for
-  for (int64_t b = 0; b < batch; ++b) {
-    int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
-    if (cache_idx == pad_slot_id) continue;
-
-    for (int64_t t = 0; t < seqlen; ++t) {
-      const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
-      scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
-      // Base of this slot in the (possibly non-contiguous) paged state
-      scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
-
-      for (int64_t d = 0; d < dim; ++d) {
-        float x_val = static_cast<float>(x_b[d * seqlen]);
-        scalar_t* sd = s_base + d * stride_s_dim;  // start of this dim's state
-        const scalar_t* w = weight_ptr + d * width;
-
-        // Accumulate in float32 for precision
-        float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
-        for (int64_t k = 0; k < state_len; ++k) {
-          acc += static_cast<float>(w[k]) *
-                 static_cast<float>(sd[k * stride_s_state]);
-        }
-        acc += static_cast<float>(w[state_len]) * x_val;
-
-        // Shift state left and append new input.
-        // Use memmove when contiguous (stride==1); element loop otherwise.
-        if (stride_s_state == 1) {
-          if (state_len > 1)
-            std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
-          if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
-        } else {
-          for (int64_t k = 0; k < state_len - 1; ++k)
-            sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
-          if (state_len > 0)
-            sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
-        }
-
-        if (do_silu) {
-          float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
-                                     : std::exp(acc) / (1.0f + std::exp(acc));
-          acc *= sigmoid;
-        }
-        out_b[d * seqlen] = static_cast<scalar_t>(acc);
-      }
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// selective_state_update
-//
-// Template parameters:
-//   state_t  - dtype of ssm_state cache (typically BFloat16)
-//   input_t  - dtype of x, B, C (typically BFloat16)
-//   out_t    - dtype of output tensor (typically BFloat16)
-//              Write directly — no float32 intermediate buffer needed.
-//
-// A, D, dt_bias are accepted as const float* (they are always float32
-// model parameters in Mamba2).  This eliminates the per-call float32→BF16
-// conversion and the .contiguous() materialisation of the broadcast-expand.
-//
-// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
-// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
-// ---------------------------------------------------------------------------
-template <typename state_t, typename input_t, typename out_t = float>
-inline void selective_state_update_kernel(
-    state_t* __restrict__ state_ptr, int64_t stride_state_n,
-    int64_t stride_state_h, int64_t stride_state_d,
-    const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
-    // dt: (N, nheads) — scalar per head, NOT expanded to head_dim
-    const float* __restrict__ dt_ptr, int64_t stride_dt_n,
-    // A: (nheads,) float32 — scalar per head
-    const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
-    const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
-    // D: (nheads,) float32 — scalar per head (nullptr if not used)
-    const float* __restrict__ D_ptr,
-    // z: same shape as x (optional)
-    const input_t* __restrict__ z_ptr,
-    // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
-    const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
-    int64_t stride_out_n, int64_t stride_out_h,
-    const int32_t* __restrict__ state_batch_indices,
-    const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
-    const int32_t* __restrict__ num_accepted_tokens,
-    const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
-    int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
-  using state_vec_t = vec_op::vec_t<state_t>;
-  using input_vec_t = vec_op::vec_t<input_t>;
-  constexpr int VEC_ELEM_NUM = 8;
-
-  int64_t nheads_per_group = nheads / ngroups;
-
-  for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
-    int64_t bos, seq_len;
-    if (cu_seqlens != nullptr) {
-      bos = cu_seqlens[seq_idx];
-      seq_len = cu_seqlens[seq_idx + 1] - bos;
-    } else {
-      bos = seq_idx;
-      seq_len = 1;
-    }
-
-    int64_t state_read_idx = (state_batch_indices != nullptr)
-                                 ? state_batch_indices[seq_idx]
-                                 : seq_idx;
-    if (state_read_idx == null_block_id) continue;
-
-    int64_t state_write_idx = (num_accepted_tokens == nullptr)
-                                  ? ((dst_state_batch_indices != nullptr)
-                                         ? dst_state_batch_indices[seq_idx]
-                                         : state_read_idx)
-                                  : -1;
-
-    state_t* s = state_ptr + state_read_idx * stride_state_n;
-
-    for (int64_t t = 0; t < seq_len; ++t) {
-      int64_t token_idx = bos + t;
-      const input_t* x_tok = x_ptr + token_idx * stride_x_n;
-      // dt: (N, nheads) — one float per head per token
-      const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
-      const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
-      const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
-      out_t* out_tok = out_ptr + token_idx * stride_out_n;
-
-#pragma omp parallel for
-      for (int64_t h = 0; h < nheads; ++h) {
-        int64_t g = h / nheads_per_group;
-        const input_t* x_h = x_tok + h * stride_x_h;
-        const input_t* B_g = B_tok + g * stride_BC_g;
-        const input_t* C_g = C_tok + g * stride_BC_g;
-        out_t* out_h = out_tok + h * stride_out_h;
-        state_t* s_h = s + h * stride_state_h;
-
-        // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
-        float dt_val = dt_tok[h];
-        if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
-        if (dt_softplus) {
-          dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
-        }
-        const float A_val = A_ptr[h];  // scalar: same for all dim, dstate
-        const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
-
-        const input_t* z_h =
-            (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
-                               : nullptr;
-
-        vec_op::FP32Vec8 dt_vec(dt_val);
-        // dA = exp(A * dt): A and dt are SCALARS per head, so compute once
-        // and broadcast.  This saves 7 redundant std::exp() calls that
-        // FP32Vec8::exp() would otherwise make on the broadcast vector.
-        const float dA_scalar = std::exp(A_val * dt_val);
-        vec_op::FP32Vec8 dA(dA_scalar);  // broadcast
-
-        for (int64_t d = 0; d < dim; ++d) {
-          float x_val = static_cast<float>(x_h[d]);
-
-          vec_op::FP32Vec8 out_vec(0.0f);
-          state_t* s_hd = s_h + d * stride_state_d;
-          const input_t* B_g_base = B_g;
-          const input_t* C_g_base = C_g;
-
-          vec_op::FP32Vec8 x_vec(x_val);
-          // dBx = B * x * dt — same dA for all dstate (A is scalar)
-          // s_new = s * dA + B * x * dt
-
-          int64_t n = 0;
-          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
-            vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
-            vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
-            vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
-
-            vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
-            vec_op::FP32Vec8 s_new = s_v * dA + dBx;
-
-            state_vec_t(s_new).save(s_hd + n);
-            out_vec = out_vec + s_new * C_v;
-          }
-
-          float out_val = out_vec.reduce_sum();
-          for (; n < dstate; ++n) {
-            // Reuse dA_scalar computed once per head — no exp() re-call
-            float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
-            float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
-            s_hd[n] = static_cast<state_t>(s_new);
-            out_val += s_new * static_cast<float>(C_g[n]);
-          }
-
-          if (D_ptr != nullptr) out_val += x_val * D_val;
-          if (z_h != nullptr) {
-            float z_val = static_cast<float>(z_h[d]);
-            float sigmoid = (z_val >= 0)
-                                ? 1.0f / (1.0f + std::exp(-z_val))
-                                : std::exp(z_val) / (1.0f + std::exp(z_val));
-            out_val *= z_val * sigmoid;
-          }
-          out_h[d] = static_cast<out_t>(out_val);
-        }
-      }
-
-      if (num_accepted_tokens != nullptr &&
-          dst_state_batch_indices != nullptr) {
-        int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
-        if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
-          state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
-          std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
-        }
-      }
-    }
-
-    if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
-        state_write_idx != state_read_idx) {
-      state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
-      std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// mamba_chunk_scan_fwd
-//
-// Prefill SSM recurrence for Mamba2 / SSD models.
-//
-// Key difference from selective_state_update_kernel (decode path):
-//   - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
-//     Each thread owns a (batch, head) slice and runs the entire token
-//     sequence without any per-token OpenMP synchronisation overhead.
-//     For seqlen=256, this eliminates 256 thread-barrier launches per batch.
-//
-// `dt` arrives already processed (float32, after bias + softplus + clamp)
-// to keep this kernel simple. Preprocessing is done in the Python wrapper.
-//
-// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
-// tensor, pre-initialised by the caller (zero or from initial_states).
-// Each (b, h) slice is private to exactly one thread via collapse(2), so
-// there are no write conflicts.
-//
-// D is treated as a scalar per head ([nheads] float32).
-// ---------------------------------------------------------------------------
-template <typename input_t>
-inline void mamba_chunk_scan_fwd_kernel(
-    float* __restrict__ states_ptr,     // [batch, nheads, headdim, dstate] f32
-    const input_t* __restrict__ x_ptr,  // [seqlen, nheads, headdim]
-    const float* __restrict__ dt_ptr,   // [seqlen, nheads] f32 (preprocessed)
-    const float* __restrict__ A_ptr,    // [nheads] f32
-    const input_t* __restrict__ B_ptr,  // [seqlen, ngroups, dstate]
-    const input_t* __restrict__ C_ptr,  // [seqlen, ngroups, dstate]
-    const float* __restrict__ D_ptr,    // [nheads] f32 (nullable)
-    const input_t* __restrict__ z_ptr,  // [seqlen, nheads, headdim] (nullable)
-    input_t* __restrict__ out_ptr,      // [seqlen, nheads, headdim]
-    const int32_t* __restrict__ cu_seqlens,  // [batch+1] int32
-    int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
-    int64_t dstate) {
-  using input_vec_t = vec_op::vec_t<input_t>;
-  constexpr int VEC_ELEM_NUM = 8;
-
-  const int64_t nheads_per_group = nheads / ngroups;
-  // states layout: [batch, nheads, headdim, dstate] contiguous (caller
-  // guarantee)
-  const int64_t stride_s_b = nheads * headdim * dstate;
-  const int64_t stride_s_h = headdim * dstate;
-  // stride_s_d = dstate, stride_s_n = 1
-
-#pragma omp parallel for collapse(2) schedule(static)
-  for (int64_t b = 0; b < batch; ++b) {
-    for (int64_t h = 0; h < nheads; ++h) {
-      const int64_t seq_start = cu_seqlens[b];
-      const int64_t seq_end = cu_seqlens[b + 1];
-      const int64_t g = h / nheads_per_group;
-
-      const float A_val = A_ptr[h];
-      const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
-
-      // Working state slice: states[b, h, :, :] — float32, headdim * dstate.
-      // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
-      float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
-
-      for (int64_t t = seq_start; t < seq_end; ++t) {
-        const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
-        const float* dt_h = dt_ptr + t * nheads + h;
-        const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
-        const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
-        const input_t* z_h = (z_ptr != nullptr)
-                                 ? z_ptr + t * nheads * headdim + h * headdim
-                                 : nullptr;
-        input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
-
-        const float dt_val = *dt_h;
-        const float dA_val = std::exp(A_val * dt_val);
-        const vec_op::FP32Vec8 dA_vec(dA_val);  // broadcast scalar
-        const vec_op::FP32Vec8 dt_vec(dt_val);
-
-        for (int64_t d = 0; d < headdim; ++d) {
-          const float x_val = static_cast<float>(x_h[d]);
-          float* s_bhd = s_bh + d * dstate;  // [dstate] contiguous float32
-
-          // Vectorised SSM update + readout over dstate:
-          //   s_new = s * dA + x * dt * B
-          //   y    += s_new * C
-          int64_t n = 0;
-          vec_op::FP32Vec8 y_vec(0.0f);
-          const vec_op::FP32Vec8 x_vec(x_val);
-
-          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
-            const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
-            const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
-            const vec_op::FP32Vec8 s_v(s_bhd + n);
-
-            const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
-            s_new.save(s_bhd + n);
-            y_vec = y_vec + s_new * C_v;
-          }
-
-          float y_val = y_vec.reduce_sum();
-
-          // Scalar tail for remaining dstate elements
-          for (; n < dstate; ++n) {
-            const float B_n = static_cast<float>(B_g[n]);
-            const float C_n = static_cast<float>(C_g[n]);
-            const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
-            s_bhd[n] = s_new;
-            y_val += s_new * C_n;
-          }
-
-          // D skip connection (scalar per head)
-          if (D_ptr != nullptr) y_val += x_val * D_val;
-
-          // z gating: out = y * z * sigmoid(z)  (SiLU)
-          if (z_h != nullptr) {
-            const float z_val = static_cast<float>(z_h[d]);
-            const float sigmoid =
-                (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
-                                : std::exp(z_val) / (1.0f + std::exp(z_val));
-            y_val *= z_val * sigmoid;
-          }
-
-          out_h[d] = static_cast<input_t>(y_val);
-        }
-      }
-    }
-  }
-}
-
-}  // namespace mamba_cpu
@@ -31,9 +31,6 @@ class MicroGemm {
  }
 };

-template <cpu_utils::ISA isa, typename scalar_t>
-class MicroGemmINT8;
-
 template <int32_t n_size, typename scalar_t>
 FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
                                   scalar_t* __restrict__ d_ptr,
@@ -1,424 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP
-#define CPU_MICRO_GEMM_INT8_NEON_HPP
-
-#include <algorithm>
-#include <cstdint>
-
-#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
-
-#include <arm_bf16.h>
-#include <arm_neon.h>
-#include <c10/util/BFloat16.h>
-#include <c10/util/Exception.h>
-#include <c10/util/Half.h>
-
-namespace cpu_micro_gemm {
-
-namespace neon_smmla {
-
-constexpr int32_t K = 8;
-constexpr int32_t Cols = 2;
-constexpr int32_t TileSize = K * Cols;
-
-FORCE_INLINE float32x4x2_t load_as_f32(const float* input) {
-  float32x4x2_t result;
-  result.val[0] = vld1q_f32(input);
-  result.val[1] = vld1q_f32(input + 4);
-  return result;
-}
-
-FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) {
-  const auto input_vec = vld1q_f16(reinterpret_cast<const float16_t*>(input));
-  float32x4x2_t result;
-  result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec));
-  result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec));
-  return result;
-}
-
-FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) {
-  const auto input_vec = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(input));
-  float32x4x2_t result;
-  result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec));
-  result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec));
-  return result;
-}
-
-FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01,
-                                    const int32x4_t acc23,
-                                    const int32x4_t acc45,
-                                    const int32x4_t acc67,
-                                    int32_t* __restrict__ c_ptr,
-                                    const int64_t ldc, const int32_t m_rows) {
-  if (m_rows == 0) {
-    return;
-  }
-
-  vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23)));
-  vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67)));
-
-  if (m_rows == 2) {
-    vst1q_s32(c_ptr + ldc,
-              vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23)));
-    vst1q_s32(c_ptr + ldc + 4,
-              vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67)));
-  }
-}
-
-FORCE_INLINE void gemm_micro_smmla_8x8_packed_a(
-    const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
-    int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
-    const int64_t ldc) {
-  const int32x4_t zero = vdupq_n_s32(0);
-  int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
-  int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
-  int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero;
-  int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero;
-
-  const int8_t* __restrict__ a_tile = a_packed;
-  const int8_t* __restrict__ b_tile = b_packed;
-
-#pragma GCC unroll 8
-  for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
-    const int8x16_t a_tile01 = vld1q_s8(a_tile);
-    const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
-    const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize);
-    const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize);
-
-    const int8x16_t b_tile01 = vld1q_s8(b_tile);
-    const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize);
-    const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize);
-    const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize);
-
-    acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
-    acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
-    acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01);
-    acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01);
-
-    acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
-    acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
-    acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23);
-    acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23);
-
-    acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
-    acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
-    acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45);
-    acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45);
-
-    acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
-    acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
-    acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67);
-    acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67);
-
-    a_tile += 4 * TileSize;
-    b_tile += 4 * TileSize;
-  }
-
-  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_smmla_4x16_packed_a(
-    const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
-    int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
-    const int64_t b_n_group_stride, const int64_t ldc) {
-  const int32_t m_rows_01 = std::min(2, m);
-  const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
-  const int32x4_t zero = vdupq_n_s32(0);
-
-  int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
-  int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
-  int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero,
-            acc011415 = zero;
-  int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero,
-            acc231415 = zero;
-
-  const int8_t* __restrict__ a_tile = a_packed;
-  // note: b packs 8 panels contiguously, so we need 2 b_tile ptrs
-  // for the 4x16 microkernel
-  const int8_t* __restrict__ b_tile0 = b_packed;
-  const int8_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 int8x16_t a_tile01 = vld1q_s8(a_tile);
-    const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
-    const int8x16_t b_tile01 = vld1q_s8(b_tile0);
-    const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize);
-    const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize);
-    const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize);
-    const int8x16_t b_tile89 = vld1q_s8(b_tile1);
-    const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize);
-    const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize);
-    const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize);
-
-    acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
-    acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
-    acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
-    acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
-
-    acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
-    acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
-    acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
-    acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
-
-    acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89);
-    acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89);
-    acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011);
-    acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011);
-
-    acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213);
-    acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213);
-    acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415);
-    acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415);
-
-    a_tile += 2 * TileSize;
-    b_tile0 += 4 * TileSize;
-    b_tile1 += 4 * TileSize;
-  }
-
-  // rows 0-1, columns 0-7
-  store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
-  // rows 0-1, columns 8-15
-  store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
-                    m_rows_01);
-  // rows 2-3, columns 0-7
-  store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
-                    m_rows_23);
-  // rows 2-3, columns 8-15
-  store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
-                    c_ptr + 2 * ldc + 8, ldc, m_rows_23);
-}
-
-}  // namespace neon_smmla
-
-template <typename scalar_t>
-class MicroGemmINT8<cpu_utils::ISA::NEON, scalar_t> {
- public:
-  static constexpr int32_t K = neon_smmla::K;
-  static constexpr int32_t Mr = 8;
-  static constexpr int32_t Nr = 8;
-  static constexpr int32_t NrGemv = 16;
-  static constexpr int32_t MaxMSize = 8;
-  static constexpr int32_t NSize = 32;
-  static constexpr int32_t WeightOCGroupSize = Nr;
-  static_assert(MaxMSize % Mr == 0);
-
-  static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
-                                        float& scale, const int32_t size) {
-    TORCH_CHECK_EQ(size % K, 0);
-    float32x4_t max_vec = vdupq_n_f32(0.0f);
-
-    for (int32_t i = 0; i < size; i += K) {
-      const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
-      max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0]));
-      max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1]));
-    }
-
-    const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f);
-    scale = abs_max / 127.0f;
-    const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max);
-
-    for (int32_t i = 0; i < size; i += K) {
-      const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
-      const int32x4_t output_low =
-          vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec));
-      const int32x4_t output_high =
-          vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec));
-      const int16x8_t output_s16 =
-          vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high));
-      vst1_s8(output + i, vqmovn_s16(output_s16));
-    }
-  }
-
-  // with current code, fusing this into the gemm micro kernel didn't move the
-  // needle
-  static FORCE_INLINE void dequantize_tile(
-      int32_t* input, float* output, const float* __restrict__ input_scales,
-      const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
-      const int32_t stride) {
-    TORCH_CHECK_EQ(n % 4, 0);
-    for (int32_t m_idx = 0; m_idx < m; ++m_idx) {
-      const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]);
-      for (int32_t n_idx = 0; n_idx < n; n_idx += 4) {
-        const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx);
-        const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx);
-        const float32x4_t output_vec =
-            vmulq_f32(vcvtq_f32_s32(input_vec),
-                      vmulq_f32(input_scale_vec, weight_scale_vec));
-        vst1q_f32(output + m_idx * stride + n_idx, output_vec);
-      }
-    }
-  }
-
-  // physical layout [
-  //  M / (8 or 4); Mr is 8 or 4
-  //  K / 8; K for smmla is 8
-  //  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 int8_t* const* __restrict__ rows,
-                                   int8_t* __restrict__ a_packed,
-                                   const int32_t m, const int32_t k) {
-    TORCH_CHECK(m > 0 && m <= MaxMSize);
-    TORCH_CHECK(k % K == 0);
-    const int8x8_t zero = vdup_n_s8(0);
-
-    for (int32_t row_base = 0; row_base < m; row_base += Mr) {
-      const int32_t panel_m = std::min(Mr, m - row_base);
-      const int8_t* const* panel_rows = rows + row_base;
-      int8_t* __restrict__ out = a_packed + row_base * k;
-
-      // fast path for full 8-row panels (fast path for 4-row panels didn't move
-      // the needle)
-      if (panel_m == Mr) {
-        const int8_t* __restrict__ row0 = panel_rows[0];
-        const int8_t* __restrict__ row1 = panel_rows[1];
-        const int8_t* __restrict__ row2 = panel_rows[2];
-        const int8_t* __restrict__ row3 = panel_rows[3];
-        const int8_t* __restrict__ row4 = panel_rows[4];
-        const int8_t* __restrict__ row5 = panel_rows[5];
-        const int8_t* __restrict__ row6 = panel_rows[6];
-        const int8_t* __restrict__ row7 = panel_rows[7];
-        int32_t k_idx = 0;
-        for (; k_idx + 2 * K <= k; k_idx += 2 * K) {
-          int8_t* __restrict__ block0 = out;
-          int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize;
-
-          int8x16_t a0 = vld1q_s8(row0 + k_idx);
-          int8x16_t a1 = vld1q_s8(row1 + k_idx);
-          vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
-          vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
-
-          a0 = vld1q_s8(row2 + k_idx);
-          a1 = vld1q_s8(row3 + k_idx);
-          vst1q_s8(block0 + neon_smmla::TileSize,
-                   vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
-          vst1q_s8(block1 + neon_smmla::TileSize,
-                   vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
-
-          a0 = vld1q_s8(row4 + k_idx);
-          a1 = vld1q_s8(row5 + k_idx);
-          vst1q_s8(block0 + 2 * neon_smmla::TileSize,
-                   vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
-          vst1q_s8(block1 + 2 * neon_smmla::TileSize,
-                   vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
-
-          a0 = vld1q_s8(row6 + k_idx);
-          a1 = vld1q_s8(row7 + k_idx);
-          vst1q_s8(block0 + 3 * neon_smmla::TileSize,
-                   vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
-          vst1q_s8(block1 + 3 * neon_smmla::TileSize,
-                   vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
-
-          out += 8 * neon_smmla::TileSize;
-        }
-
-        for (; k_idx < k; k_idx += K) {
-          int8x8_t a0 = vld1_s8(row0 + k_idx);
-          int8x8_t a1 = vld1_s8(row1 + k_idx);
-          vst1q_s8(out, vcombine_s8(a0, a1));
-
-          a0 = vld1_s8(row2 + k_idx);
-          a1 = vld1_s8(row3 + k_idx);
-          vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1));
-
-          a0 = vld1_s8(row4 + k_idx);
-          a1 = vld1_s8(row5 + k_idx);
-          vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1));
-
-          a0 = vld1_s8(row6 + k_idx);
-          a1 = vld1_s8(row7 + k_idx);
-          vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1));
-
-          out += 4 * neon_smmla::TileSize;
-        }
-        continue;
-      }
-
-      const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2;
-      for (int32_t k_idx = 0; k_idx < k; k_idx += K) {
-        for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) {
-          const int32_t row_idx = pair_idx * 2;
-          const int8x8_t row0 =
-              (row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero;
-          const int8x8_t row1 = (row_idx + 1 < panel_m)
-                                    ? vld1_s8(panel_rows[row_idx + 1] + k_idx)
-                                    : zero;
-          vst1q_s8(out, vcombine_s8(row0, row1));
-          out += neon_smmla::TileSize;
-        }
-      }
-    }
-  }
-
-  // physical layout [
-  //  N / 8; Nr is 8
-  //  K / 8; K for smmla is 8
-  //  4,   ; 4 col-pairs for each 8 cols
-  //  2,   ; col-pair is 2 cols
-  //  4    ; 4 elements per col
-  // ]
-  static void pack_weight(const int8_t* __restrict__ weight,
-                          int8_t* __restrict__ packed_weight,
-                          const int32_t output_size, const int32_t input_size) {
-    TORCH_CHECK(output_size % NSize == 0);
-    TORCH_CHECK(input_size % K == 0);
-
-    for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
-      int8_t* __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 += neon_smmla::Cols) {
-          const int8_t* __restrict__ row0 =
-              weight + (o_idx + pair_idx) * input_size + k_idx;
-          const int8_t* __restrict__ row1 = row0 + input_size;
-          vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1)));
-          dst += neon_smmla::TileSize;
-        }
-      }
-    }
-  }
-
-  void gemm(const int8_t* __restrict__ a_packed,
-            const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
-            const int32_t m, const int32_t k, const int64_t b_n_group_stride,
-            const int64_t ldc) const {
-    TORCH_CHECK(m > 0 && m <= MaxMSize);
-    TORCH_CHECK(k % K == 0);
-
-    for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
-      const int8_t* __restrict__ b_panel = b_packed + 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 int8_t* __restrict__ a_panel = a_packed + row_base * k;
-        int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx;
-
-        if (panel_m <= 4) {
-          neon_smmla::gemm_micro_smmla_4x16_packed_a(
-              a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc);
-        } else {
-          neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel,
-                                                    panel_m, k, ldc);
-          neon_smmla::gemm_micro_smmla_8x8_packed_a(
-              a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k,
-              ldc);
-        }
-      }
-    }
-  }
-};
-
-}  // namespace cpu_micro_gemm
-
-#endif
@@ -1,6 +1,3 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
 #ifndef CPU_MICRO_GEMM_NEON_HPP
 #define CPU_MICRO_GEMM_NEON_HPP

@@ -19,6 +16,9 @@ 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) {
@@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
    acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);

    a_tile += 4 * TileSize;
-    b_tile += 4 * TileSize;
+    b_tile += Nr * K;
  }

  store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
@@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
    acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);

    a_tile += 2 * TileSize;
-    b_tile0 += 4 * TileSize;
-    b_tile1 += 4 * TileSize;
+    b_tile0 += Nr * K;
+    b_tile1 += Nr * K;
  }

  store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
@@ -223,9 +223,6 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
 template <typename scalar_t>
 class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
 public:
-  static constexpr int32_t Mr = 8;
-  static constexpr int32_t Nr = 8;
-  static constexpr int32_t NrGemv = 16;
  static constexpr int32_t MaxMSize = 8;
  static constexpr int32_t NSize = 32;
  static constexpr int32_t WeightOCGroupSize = Nr;
@@ -249,9 +246,6 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
 public:
  using scalar_t = c10::BFloat16;

-  static constexpr int32_t Mr = 8;
-  static constexpr int32_t Nr = 8;
-  static constexpr int32_t NrGemv = 16;
  static constexpr int32_t MaxMSize = 8;
  static constexpr int32_t NSize = 32;
  static constexpr int32_t WeightOCGroupSize = Nr;
@@ -259,7 +253,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {

 public:
  // physical layout [
-  //  M / (8 or 4); Mr is 8 or 4
+  //  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
@@ -445,7 +439,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
    (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 += NrGemv) {
+    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;

@@ -451,90 +451,6 @@ void causal_conv1d_update_kernel_impl(
  });
 }

-template <typename scalar_t>
-void causal_conv1d_update_multi_kernel_impl(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    scalar_t* __restrict__ conv_states,
-    const scalar_t* __restrict__ weight,
-    const scalar_t* __restrict__ bias,
-    const int32_t* __restrict__ num_accepted_tokens,
-    const int32_t* __restrict__ conv_indices,
-    bool silu_activation,
-    int64_t batch,
-    int64_t dim,
-    int64_t seqlen,
-    int64_t width,
-    int64_t state_len,
-    int64_t conv_state_slot_stride) {
-  constexpr int64_t BLOCK_N = block_size_n() * 2;
-  const int64_t NB = div_up(dim, BLOCK_N);
-
-  AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
-    at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
-      int64_t bs{0}, nb{0};
-      data_index_init(begin, bs, batch, nb, NB);
-
-      for (int64_t i = begin; i < end; ++i) {
-        const int64_t nb_start = nb * BLOCK_N;
-        const int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
-        const int32_t conv_state_index = conv_indices[bs];
-        const int32_t history_offset = num_accepted_tokens[bs] - 1;
-
-        switch (width << 4 | nb_size >> 4) {
-          case 0x42:
-            tinygemm_kernel<scalar_t, 4, 32, has_bias, has_silu>::apply(
-                input + bs * seqlen * dim + nb_start,
-                weight + nb_start * width,
-                out + bs * seqlen * dim + nb_start,
-                has_bias ? bias + nb_start : nullptr,
-                conv_states + conv_state_index * conv_state_slot_stride +
-                    history_offset * dim + nb_start,
-                true,
-                seqlen,
-                dim,
-                true);
-            break;
-          case 0x44:
-            tinygemm_kernel<scalar_t, 4, 64, has_bias, has_silu>::apply(
-                input + bs * seqlen * dim + nb_start,
-                weight + nb_start * width,
-                out + bs * seqlen * dim + nb_start,
-                has_bias ? bias + nb_start : nullptr,
-                conv_states + conv_state_index * conv_state_slot_stride +
-                    history_offset * dim + nb_start,
-                true,
-                seqlen,
-                dim,
-                true);
-            break;
-          default:
-            TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
-        }
-
-        data_index_step(bs, batch, nb, NB);
-      }
-    });
-  });
-
-  at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t bs = begin; bs < end; ++bs) {
-      const int32_t conv_state_index = conv_indices[bs];
-      const int32_t num_accepted = num_accepted_tokens[bs];
-      scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride;
-
-      std::memmove(
-          state,
-          state + num_accepted * dim,
-          (state_len - seqlen) * dim * sizeof(scalar_t));
-      std::memcpy(
-          state + (state_len - seqlen) * dim,
-          input + bs * seqlen * dim,
-          seqlen * dim * sizeof(scalar_t));
-    }
-  });
-}
-
 }  // anonymous namespace

 // from [dim, width] or [N, K]
@@ -629,7 +545,7 @@ at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t n
 //   query_start_loc: (batch + 1) int32
 //   cache_indices: (batch)  int32
 //   has_initial_state: (batch) bool
-//   conv_states: (..., dim, state_len) itype, where state_len >= width - 1
+//   conv_states: (..., dim, width - 1) itype
 //   activation: either None or "silu" or "swish"
 //   pad_slot_id: int
 //
@@ -670,14 +586,11 @@ at::Tensor causal_conv1d_fwd_cpu(
    CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
    CHECK_GE(padded_batch, batch);
    CHECK_EQ(conv_states_val.size(1), dim);
-    const int64_t state_len = conv_states_val.size(2);
-    CHECK_GE(state_len, width - 1);
+    CHECK_EQ(conv_states_val.size(2), width - 1);

    // adjust `conv_states` to be contiguous on `dim`
    // should happen only once
    if (conv_states_val.stride(-2) != 1) {
-      TORCH_CHECK(state_len == width - 1,
-          "causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim.");
      auto conv_states_copy = conv_states_val.clone();
      conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
      conv_states_val.copy_(conv_states_copy);
@@ -738,14 +651,14 @@ at::Tensor causal_conv1d_fwd_cpu(

 // API aligned with GPUs
 //
-//   x: (batch, dim) or (batch, seqlen, dim)
+//   x: (batch, dim) or (batch, dim, seqlen)
 //   conv_state: (..., dim, state_len), where state_len >= width - 1
 //   weight: (dim, width)
 //   bias: (dim,)
-//   num_accepted_tokens: (batch,), dtype int32.
+//   cache_seqlens: (batch,), dtype int32.
 //   conv_state_indices: (batch,), dtype int32
 //   pad_slot_id: int
-//   out: (batch, dim) or (batch, seqlen, dim)
+//   out: (batch, dim) or (batch, dim, seqlen)
 //
 at::Tensor causal_conv1d_update_cpu(
    const at::Tensor& x,
@@ -753,7 +666,7 @@ at::Tensor causal_conv1d_update_cpu(
    const at::Tensor& weight,
    const std::optional<at::Tensor>& bias,
    bool silu_activation,
-    const std::optional<at::Tensor>& num_accepted_tokens,
+    const std::optional<at::Tensor>& cache_seqlens,
    const std::optional<at::Tensor>& conv_state_indices,
    int64_t pad_slot_id,
    bool is_vnni) {
@@ -761,13 +674,13 @@ at::Tensor causal_conv1d_update_cpu(
  CHECK_CONTIGUOUS(weight);
  auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);

-  TORCH_CHECK(
-      x.dim() == 2 || x.dim() == 3,
-      "causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");
+  // TODO: add multi-token prediction support
+  TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
+  TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");

  int64_t batch = x.size(0);
-  int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);
-  int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);
+  int64_t dim = x.size(1);
+  int64_t seqlen = 1;
  int64_t width = weight.size(-1);

  const auto scalar_type = x.scalar_type();
@@ -777,84 +690,10 @@ at::Tensor causal_conv1d_update_cpu(

  CHECK_EQ(conv_states.scalar_type(), scalar_type);
  CHECK_EQ(conv_states.size(1), dim);
-  const int64_t state_len = conv_states.size(2);
-  CHECK_GE(state_len, width - 1);
-
-  if (x.dim() == 3) {
-    TORCH_CHECK(
-        num_accepted_tokens.has_value(),
-        "causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x.");
-    TORCH_CHECK(
-        conv_state_indices.has_value(),
-        "causal_conv1d_update_cpu: conv_state_indices is required for 3D x.");
-    CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt);
-    TORCH_CHECK(
-        width == 4,
-        "causal_conv1d_update_cpu: support only width of 4 for 3D x.");
-    TORCH_CHECK(
-        seqlen > 0,
-        "causal_conv1d_update_cpu: expect non-empty sequence for 3D x.");
-    TORCH_CHECK(
-        state_len >= seqlen,
-        "causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x.");
-    TORCH_CHECK(
-        conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim,
-        "causal_conv1d_update_cpu: 3D x requires SD conv_states layout.");
-
-    const int32_t* accepted_counts =
-        num_accepted_tokens.value().data_ptr<int32_t>();
-    const int32_t* indices = conv_state_indices.value().data_ptr<int32_t>();
-    const int64_t num_slots = conv_states.size(0);
-    for (int64_t bs = 0; bs < batch; ++bs) {
-      const int32_t num_accepted = accepted_counts[bs];
-      const int32_t conv_state_index = indices[bs];
-      TORCH_CHECK(
-          conv_state_index != pad_slot_id,
-          "causal_conv1d_update_cpu: 3D x does not support pad slots.");
-      TORCH_CHECK(
-          conv_state_index >= 0 && conv_state_index < num_slots,
-          "causal_conv1d_update_cpu: conv_state_indices out of range.");
-      TORCH_CHECK(
-          num_accepted >= 1 && num_accepted <= seqlen,
-          "causal_conv1d_update_cpu: num_accepted_tokens must be in [1, "
-          "seqlen].");
-      TORCH_CHECK(
-          num_accepted - 1 + width - 1 <= state_len,
-          "causal_conv1d_update_cpu: history window exceeds conv_states.");
-    }
-
-    int64_t conv_state_slot_stride = conv_states.stride(0);
-    at::Tensor out = at::empty_like(x);
-    AT_DISPATCH_REDUCED_FLOATING_TYPES(
-        scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] {
-          causal_conv1d_update_multi_kernel_impl<scalar_t>(
-              out.data_ptr<scalar_t>(),
-              x.data_ptr<scalar_t>(),
-              conv_states.data_ptr<scalar_t>(),
-              packed_w.data_ptr<scalar_t>(),
-              conditional_data_ptr<scalar_t>(bias),
-              accepted_counts,
-              indices,
-              silu_activation,
-              batch,
-              dim,
-              seqlen,
-              width,
-              state_len,
-              conv_state_slot_stride);
-        });
-    return out;
-  }
-
-  TORCH_CHECK(
-      !num_accepted_tokens.has_value(),
-      "causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D "
-      "x.");
+  CHECK_EQ(conv_states.size(2), width - 1);

  // adjust `conv_states` to be contiguous on `dim`
  if (conv_states.stride(-2) != 1) {
-    TORCH_CHECK(state_len == width - 1,
-        "causal_conv1d_update_cpu: wide conv_states must be contiguous on dim.");
    int64_t num_cache_lines = conv_states.size(0);
    auto conv_states_copy = conv_states.clone();
    conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
@@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu(
 at::Tensor causal_conv1d_update_cpu(
    const at::Tensor& x, const at::Tensor& conv_states,
    const at::Tensor& weight, const std::optional<at::Tensor>& bias,
-    bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,
+    bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
    const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
    bool is_vnni);

@@ -163,8 +163,7 @@ torch::Tensor get_scheduler_metadata(
    const torch::Tensor& query_start_loc, const bool casual,
    const int64_t window_size, const std::string& isa_hint,
    const bool enable_kv_split,
-    const std::optional<torch::Tensor>& dynamic_causal,
-    const std::string& kv_cache_dtype);
+    const std::optional<torch::Tensor>& dynamic_causal);

 void cpu_attn_reshape_and_cache(const torch::Tensor& key,
                                const torch::Tensor& value,
@@ -208,52 +207,12 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
                   const torch::Tensor& topk_id, const bool skip_weighted,
                   const std::string& act, const std::string& isa);

-void prepack_moe_weight_int8(const torch::Tensor& weight,
-                             torch::Tensor& packed_weight,
-                             const std::string& isa);
-
-void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
-                        const torch::Tensor& w13, const torch::Tensor& w2,
-                        const torch::Tensor& w13_scale,
-                        const torch::Tensor& w2_scale,
-                        const std::optional<torch::Tensor>& w13_bias,
-                        const std::optional<torch::Tensor>& w2_bias,
-                        const torch::Tensor& topk_weights,
-                        const torch::Tensor& topk_id, const bool skip_weighted,
-                        const std::string& act, const std::string& isa);
-
 void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                      const torch::Tensor positions,
                                      const torch::Tensor block_table,
                                      torch::Tensor slot_mapping,
                                      const int64_t block_size);

-at::Tensor causal_conv1d_update_cpu_impl(
-    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
-    const c10::optional<at::Tensor>& bias,
-    const c10::optional<std::string>& activation,
-    const c10::optional<at::Tensor>& conv_state_indices,
-    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
-
-void selective_state_update_cpu_impl(
-    at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
-    const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
-    const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
-    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
-    const c10::optional<at::Tensor>& state_batch_indices,
-    const c10::optional<at::Tensor>& dst_state_batch_indices,
-    int64_t null_block_id, at::Tensor& out,
-    const c10::optional<at::Tensor>& num_accepted_tokens,
-    const c10::optional<at::Tensor>& cu_seqlens);
-
-void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
-                                   const at::Tensor& x, const at::Tensor& dt,
-                                   const at::Tensor& A, const at::Tensor& B,
-                                   const at::Tensor& C,
-                                   const c10::optional<at::Tensor>& D,
-                                   const c10::optional<at::Tensor>& z,
-                                   const at::Tensor& cu_seqlens);
-
 void init_cpu_memory_env(std::vector<int64_t> node_ids);

 namespace cpu_utils {
@@ -517,8 +476,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
      "weight, Tensor? bias, bool silu_activation,"
-      "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "
-      "pad_slot_id, "
+      "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
      "bool is_vnni) -> Tensor");
  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
 #endif
@@ -578,8 +536,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
      "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
      "query_start_loc, bool casual, int window_size, str isa_hint, bool "
-      "enable_kv_split, Tensor? dynamic_causal, "
-      "str kv_cache_dtype=\"auto\") -> Tensor",
+      "enable_kv_split, Tensor? dynamic_causal) -> Tensor",
      &get_scheduler_metadata);
  ops.def(
      "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -613,7 +570,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif

  // fused moe
-#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
+#if defined(__AVX512F__) || \
+    (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
  ops.def(
      "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
      "-> ()");
@@ -624,22 +582,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) &&
-        // !defined(__APPLE__))
-#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
-    !defined(__APPLE__)
-  ops.def(
-      "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
-      "str isa) -> ()");
-  ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8);
-  ops.def(
-      "cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, "
-      "Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, "
-      "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
-      "skip_weighted, str act, str isa) -> ()");
-  ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
-#endif  // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
-        // !defined(__APPLE__)
+#endif
  ops.def(
      "mla_decode_kvcache("
      "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -652,30 +595,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "block_size) -> ()",
      &compute_slot_mapping_kernel_impl);

-  // Mamba CPU kernels
-  ops.def(
-      "causal_conv1d_update_cpu_vec("
-      "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
-      "Tensor? bias, str? activation, Tensor? conv_state_indices, "
-      "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
-      &causal_conv1d_update_cpu_impl);
-
-  ops.def(
-      "selective_state_update_cpu("
-      "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
-      "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
-      "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
-      "SymInt null_block_id, Tensor(a13!) out, "
-      "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
-      &selective_state_update_cpu_impl);
-
-  ops.def(
-      "mamba_chunk_scan_fwd_cpu("
-      "Tensor(a0!) out, Tensor(a1!) final_states, "
-      "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
-      "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
-      &mamba_chunk_scan_fwd_cpu_impl);
-
  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);

  // Speculative decoding kernels
@@ -1,327 +0,0 @@
-#pragma once
-
-#include "custom_collective_common.cuh"
-
-namespace vllm {
-
-constexpr int kMnnvlLamportAgThreads = 128;
-constexpr int kMnnvlLamportRsThreads = 256;
-constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
-
-using CopyPack = array_t<uint64_t, 2>;
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1)
-    cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg,
-                            CopyPack* __restrict__ result, int rank,
-                            int size_per_rank) {
-  auto dp = *_dp;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  barrier_at_start<ngpus>(sg, self_sg, rank);
-#pragma unroll
-  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-    auto src = reinterpret_cast<const CopyPack*>(dp.ptrs[src_rank]);
-    auto dst = result + src_rank * size_per_rank;
-    for (int idx = tid; idx < size_per_rank; idx += stride) {
-      dst[idx] = src[idx];
-    }
-  }
-  barrier_at_end<ngpus, true>(sg, self_sg, rank);
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(512, 1)
-    cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg,
-                                T* __restrict__ result, int rank,
-                                int size_per_rank) {
-  using P = typename packed_t<T>::P;
-  using A = typename packed_t<T>::A;
-  auto dp = *_dp;
-  auto offset = rank * size_per_rank;
-  barrier_at_start<ngpus>(sg, self_sg, rank);
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank;
-       idx += gridDim.x * blockDim.x) {
-    reinterpret_cast<P*>(result)[idx] =
-        packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
-  }
-  barrier_at_end<ngpus, true>(sg, self_sg, rank);
-}
-
-template <typename P>
-union LamportPack {
-  P packed;
-  uint32_t words[sizeof(P) / sizeof(uint32_t)];
-};
-
-template <typename P>
-DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
-  static_assert(sizeof(P) == 16);
-  LamportPack<P> value;
-#if !defined(USE_ROCM)
-  asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
-               : "=r"(value.words[0]), "=r"(value.words[1]),
-                 "=r"(value.words[2]), "=r"(value.words[3])
-               : "l"(ptr)
-               : "memory");
-#else
-  const volatile uint32_t* src =
-      reinterpret_cast<const volatile uint32_t*>(ptr);
-  #pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    value.words[i] = src[i];
-  }
-#endif
-  return value;
-}
-
-template <typename P>
-DINLINE bool is_lamport_dirty(const LamportPack<P>& value) {
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    if (value.words[i] == 0x80000000U) return true;
-  }
-  return false;
-}
-
-template <typename P>
-DINLINE P lamport_sentinel() {
-  LamportPack<P> value;
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    value.words[i] = 0x80000000U;
-  }
-  return value.packed;
-}
-
-template <typename P>
-DINLINE P sanitize_lamport_payload(P packed) {
-  LamportPack<P> value{.packed = packed};
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    if (value.words[i] == 0x80000000U) value.words[i] = 0;
-  }
-  return value.packed;
-}
-
-template <typename P>
-DINLINE P wait_lamport_payload(const P* ptr) {
-  auto value = load_lamport_pack(ptr);
-  while (is_lamport_dirty(value)) value = load_lamport_pack(ptr);
-  return value.packed;
-}
-
-template <typename P, int ngpus>
-DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride,
-                                   P local_value, P (&values)[ngpus]) {
-  bool ready[ngpus];
-#pragma unroll
-  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-    ready[src_rank] = src_rank == rank;
-    if (src_rank == rank) values[src_rank] = local_value;
-  }
-
-  int remaining = ngpus - 1;
-  while (remaining != 0) {
-#pragma unroll
-    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-      if (!ready[src_rank]) {
-        auto value = load_lamport_pack(base + src_rank * rank_stride);
-        if (!is_lamport_dirty(value)) {
-          values[src_rank] = value.packed;
-          ready[src_rank] = true;
-          --remaining;
-        }
-      }
-    }
-  }
-}
-
-template <typename P, typename A, int ngpus>
-DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input,
-                                  int rank, int size_per_rank, int idx) {
-  P source_zero =
-      rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx);
-  A tmp = upcast(source_zero);
-#pragma unroll
-  for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
-    P value = src_rank == rank
-                  ? packed_input[rank * size_per_rank + idx]
-                  : wait_lamport_payload(current_local +
-                                         src_rank * size_per_rank + idx);
-    packed_assign_add(tmp, upcast(value));
-  }
-  return sanitize_lamport_payload(downcast<P>(tmp));
-}
-
-DINLINE void lamport_cta_arrive(uint32_t* counter) {
-#if !defined(USE_ROCM)
-  if (threadIdx.x < 32) {
-    asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory");
-    if (threadIdx.x == 0) {
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-      asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;"
-                   :
-                   : "l"(counter)
-                   : "memory");
-  #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-      asm volatile("red.release.global.gpu.add.u32 [%0], 1;"
-                   :
-                   : "l"(counter)
-                   : "memory");
-  #else
-      atomicAdd(counter, 1);
-  #endif
-    }
-  } else {
-    asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory");
-  }
-#else
-  __syncthreads();
-  if (threadIdx.x == 0) atomicAdd(counter, 1);
-#endif
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
-    mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input,
-                             T* __restrict__ result,
-                             T* __restrict__ multicast_buffer,
-                             uint32_t* __restrict__ epochs, int rank,
-                             int size_per_rank, int stage_size) {
-  using P = typename packed_t<T>::P;
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaGridDependencySynchronize();
-#endif
-  auto dp = *_dp;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  uint32_t epoch = epochs[0];
-  int current_stage = epoch % 3;
-  int dirty_stage = (epoch + 1) % 3;
-  int dirty_size = epochs[2 + dirty_stage];
-  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
-  auto current_local = local_buffer + current_stage * stage_size;
-  auto dirty_local = local_buffer + dirty_stage * stage_size;
-  auto current_multicast =
-      reinterpret_cast<P*>(multicast_buffer) + current_stage * stage_size;
-  auto packed_input = reinterpret_cast<const P*>(input);
-  auto packed_result = reinterpret_cast<P*>(result);
-
-  int total_size = size_per_rank * ngpus;
-  P local_value;
-  if (tid < size_per_rank) {
-    local_value = packed_input[tid];
-    current_multicast[rank * size_per_rank + tid] =
-        sanitize_lamport_payload(local_value);
-  }
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
-
-  lamport_cta_arrive(&epochs[1]);
-
-  for (int idx = tid; idx < dirty_size; idx += stride) {
-    dirty_local[idx] = lamport_sentinel<P>();
-  }
-
-  if (tid < size_per_rank) {
-#pragma unroll
-    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-      int output_idx = src_rank * size_per_rank + tid;
-      P value = src_rank == rank
-                    ? local_value
-                    : wait_lamport_payload(current_local + output_idx);
-      packed_result[output_idx] = value;
-    }
-  }
-
-  if (tid == 0) {
-    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
-    epochs[2 + current_stage] = total_size;
-    epochs[0] = epoch + 1;
-    epochs[1] = 0;
-  }
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
-    mnnvl_lamport_reduce_scatter_kernel(RankData* _dp,
-                                        const T* __restrict__ input,
-                                        T* __restrict__ result,
-                                        uint32_t* __restrict__ epochs, int rank,
-                                        int size_per_rank, int stage_size) {
-  using P = typename packed_t<T>::P;
-  using A = typename packed_t<T>::A;
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaGridDependencySynchronize();
-#endif
-  auto dp = *_dp;
-  int dst_rank = blockIdx.x % ngpus;
-  int tile = blockIdx.x / ngpus;
-  int idx = tile * blockDim.x + threadIdx.x;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  uint32_t epoch = epochs[0];
-  int current_stage = epoch % 3;
-  int dirty_stage = (epoch + 1) % 3;
-  int dirty_size = epochs[2 + dirty_stage];
-  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
-  auto current_local = local_buffer + current_stage * stage_size;
-  auto dirty_local = local_buffer + dirty_stage * stage_size;
-  auto packed_input = reinterpret_cast<const P*>(input);
-
-  if (idx < size_per_rank && dst_rank != rank) {
-    auto dst = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[dst_rank])) +
-               current_stage * stage_size + rank * size_per_rank;
-    auto src = packed_input + dst_rank * size_per_rank;
-    dst[idx] = sanitize_lamport_payload(src[idx]);
-  }
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
-
-  lamport_cta_arrive(&epochs[1]);
-
-  for (int idx = tid; idx < dirty_size; idx += stride) {
-    dirty_local[idx] = lamport_sentinel<P>();
-  }
-
-  if (idx < size_per_rank && dst_rank == rank) {
-    if constexpr (ngpus == 4) {
-      if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) {
-        reinterpret_cast<P*>(result)[idx] =
-            reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input,
-                                                 rank, size_per_rank, idx);
-      } else {
-        P values[ngpus];
-        wait_lamport_payloads<P, ngpus>(
-            current_local + idx, rank, size_per_rank,
-            packed_input[rank * size_per_rank + idx], values);
-        A tmp = upcast(values[0]);
-#pragma unroll
-        for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
-          packed_assign_add(tmp, upcast(values[src_rank]));
-        }
-        reinterpret_cast<P*>(result)[idx] =
-            sanitize_lamport_payload(downcast<P>(tmp));
-      }
-    } else {
-      reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>(
-          current_local, packed_input, rank, size_per_rank, idx);
-    }
-  }
-
-  if (tid == 0) {
-    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
-    epochs[2 + current_stage] = size_per_rank * ngpus;
-    epochs[0] = epoch + 1;
-    epochs[1] = 0;
-  }
-}
-
-}  // namespace vllm
@@ -1,8 +1,299 @@
 #pragma once

-#include "custom_collective_common.cuh"
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#if defined(USE_ROCM)
+typedef __hip_bfloat16 nv_bfloat16;
+#endif
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+#include <cstdlib>
+#include <cstring>

 namespace vllm {
+#define CUDACHECK(cmd)                                              \
+  do {                                                              \
+    cudaError_t e = cmd;                                            \
+    if (e != cudaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             cudaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+// Maximal number of blocks in allreduce kernel.
+constexpr int kMaxBlocks = 36;
+
+// Default number of blocks in allreduce kernel.
+#ifndef USE_ROCM
+const int defaultBlockLimit = 36;
+CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
+#else
+const int defaultBlockLimit = 16;
+hipPointer_attribute rangeStartAddrAttr =
+    HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
+#endif
+
+// Counter may overflow, but it's fine since unsigned int overflow is
+// well-defined behavior.
+using FlagType = uint32_t;
+
+// Two sets of peer counters are needed for two syncs: starting and ending an
+// operation. The reason is that it's possible for peer GPU block to arrive at
+// the second sync point while the current GPU block haven't passed the first
+// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
+// waiting for counter. We use alternating counter array to avoid this
+// possibility.
+struct Signal {
+  alignas(128) FlagType start[kMaxBlocks][8];
+  alignas(128) FlagType end[kMaxBlocks][8];
+  alignas(128) FlagType _flag[kMaxBlocks];  // incremental flags for each rank
+};
+
+struct __align__(16) RankData {
+  const void* ptrs[8];
+};
+
+struct __align__(16) RankSignals {
+  Signal* signals[8];
+};
+
+// like std::array, but aligned
+template <typename T, int sz>
+struct __align__(alignof(T) * sz) array_t {
+  T data[sz];
+  using type = T;
+  static constexpr int size = sz;
+};
+
+// use packed type to maximize memory efficiency
+// goal: generate ld.128 and st.128 instructions
+template <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+DINLINE T downcast_s(float val);
+template <>
+DINLINE half downcast_s(float val) {
+  return __float2half(val);
+}
+
+// scalar add functions
+// for some reason when compiling with Pytorch, the + operator for half and
+// bfloat is disabled so we call the intrinsics directly
+DINLINE half& assign_add(half& a, half b) {
+  a = __hadd(a, b);
+  return a;
+}
+DINLINE float& assign_add(float& a, float b) { return a += b; }
+
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
+template <>
+DINLINE nv_bfloat16 downcast_s(float val) {
+  return __float2bfloat16(val);
+}
+DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
+  a = __hadd(a, b);
+  return a;
+}
+#endif
+
+template <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(val.data[i]);
+    }
+    return out;
+  }
+}
+
+#if !defined(USE_ROCM)
+
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+  #else
+  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+  #endif
+}
+
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
+  FlagType flag;
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  #else
+  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  #endif
+  return flag;
+}
+
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+}
+
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
+  FlagType flag;
+  asm volatile("ld.volatile.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  return flag;
+}
+
+// This function is meant to be used as the first synchronization in the all
+// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
+// prior memory accesses. Note: volatile writes will not be reordered against
+// other volatile writes.
+template <int ngpus>
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
+                              int rank) {
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
+    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
+    // Write the expected counter value to peer and wait for correct value
+    // from peer.
+    st_flag_volatile(peer_counter_ptr, flag);
+    while (ld_flag_volatile(self_counter_ptr) != flag);
+  }
+  __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+// This function is meant to be used as the second or the final
+// synchronization barrier in the all reduce kernel. If it's the final
+// synchronization barrier, we don't need to make any visibility guarantees
+// for prior memory accesses.
+template <int ngpus, bool final_sync = false>
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
+  __syncthreads();
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
+    auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
+    // Write the expected counter value to peer and wait for correct value from
+    // peer.
+    if constexpr (!final_sync) {
+      st_flag_release(peer_counter_ptr, flag);
+      while (ld_flag_acquire(self_counter_ptr) != flag);
+    } else {
+      st_flag_volatile(peer_counter_ptr, flag);
+      while (ld_flag_volatile(self_counter_ptr) != flag);
+    }
+  }
+  if constexpr (!final_sync) __syncthreads();
+
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+#else
+
+template <int ngpus>
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
+                              int rank) {
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    // simultaneously write to the corresponding flag of all ranks.
+    // Latency = 1 p2p write
+    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
+                            flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
+    // wait until we got true from all ranks
+    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
+                                  __ATOMIC_RELAXED,
+                                  __MEMORY_SCOPE_DEVICE) < flag);
+  }
+  __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+template <int ngpus, bool final_sync = false>
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
+  __syncthreads();
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    // simultaneously write to the corresponding flag of all ranks.
+    // Latency = 1 p2p write
+    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
+                            flag,
+                            final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
+                            __MEMORY_SCOPE_SYSTEM);
+    // wait until we got true from all ranks
+    while (
+        __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
+                               final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
+                               __MEMORY_SCOPE_DEVICE) < flag);
+  }
+  if constexpr (!final_sync) __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+#endif
+
+template <typename P, int ngpus, typename A>
+DINLINE P packed_reduce(const P* ptrs[], int idx) {
+  A tmp = upcast(ptrs[0][idx]);
+#pragma unroll
+  for (int i = 1; i < ngpus; i++) {
+    packed_assign_add(tmp, upcast(ptrs[i][idx]));
+  }
+  return downcast<P>(tmp);
+}

 template <typename T, int ngpus>
 __global__ void __launch_bounds__(512, 1)
@@ -325,21 +616,6 @@ class CustomAllreduce {
 #undef KL
  }

-  void allgather(cudaStream_t stream, void* input, void* output, int size_bytes,
-                 int threads = 512, int block_limit = defaultBlockLimit);
-  template <typename T>
-  void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output,
-                               void* local_buffer, void* multicast_buffer,
-                               uint32_t* epochs, int size_bytes,
-                               int stage_size_bytes);
-  template <typename T>
-  void reduce_scatter(cudaStream_t stream, T* input, T* output, int size,
-                      int threads = 512, int block_limit = defaultBlockLimit);
-  template <typename T>
-  void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output,
-                                    void* local_buffer, uint32_t* epochs,
-                                    int size, int stage_size_bytes);
-
  ~CustomAllreduce() {
    for (auto [_, ptr] : ipc_handles_) {
      CUDACHECK(cudaIpcCloseMemHandle(ptr));
@@ -349,8 +625,8 @@ class CustomAllreduce {

 /**
 * To inspect PTX/SASS, copy paste this header file to compiler explorer and
- * add a template instantiation:
+ add a template instantiation:
 * template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
- *                                                       half *, int, int, int);
- */
-}  // namespace vllm
+ half *, int, int, int);
+*/
+}  // namespace vllm
@@ -1,332 +0,0 @@
-#pragma once
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-#if defined(USE_ROCM)
-typedef __hip_bfloat16 nv_bfloat16;
-#endif
-
-#include <iostream>
-#include <array>
-#include <limits>
-#include <map>
-#include <unordered_map>
-#include <vector>
-#include <cstdlib>
-#include <cstring>
-
-namespace vllm {
-constexpr int kMaxCustomCollectiveRanks = 16;
-
-#define CUDACHECK(cmd)                                              \
-  do {                                                              \
-    cudaError_t e = cmd;                                            \
-    if (e != cudaSuccess) {                                         \
-      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
-             cudaGetErrorString(e));                                \
-      exit(EXIT_FAILURE);                                           \
-    }                                                               \
-  } while (0)
-
-// Maximal number of blocks in allreduce kernel.
-constexpr int kMaxBlocks = 36;
-
-// Default number of blocks in allreduce kernel.
-#ifndef USE_ROCM
-inline constexpr int defaultBlockLimit = 36;
-inline CUpointer_attribute rangeStartAddrAttr =
-    CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
-#else
-inline constexpr int defaultBlockLimit = 16;
-inline hipPointer_attribute rangeStartAddrAttr =
-    HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
-#endif
-
-// Counter may overflow, but it's fine since unsigned int overflow is
-// well-defined behavior.
-using FlagType = uint32_t;
-
-// Two sets of peer counters are needed for two syncs: starting and ending an
-// operation. The reason is that it's possible for peer GPU block to arrive at
-// the second sync point while the current GPU block haven't passed the first
-// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
-// waiting for counter. We use alternating counter array to avoid this
-// possibility.
-struct Signal {
-  alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks];
-  alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks];
-  alignas(128) FlagType _flag[kMaxBlocks];  // incremental flags for each rank
-};
-
-struct __align__(16) RankData {
-  const void* ptrs[kMaxCustomCollectiveRanks];
-};
-
-struct __align__(16) RankSignals {
-  Signal* signals[kMaxCustomCollectiveRanks];
-};
-
-// like std::array, but aligned
-template <typename T, int sz>
-struct __align__(alignof(T) * sz) array_t {
-  T data[sz];
-  using type = T;
-  static constexpr int size = sz;
-};
-
-// use packed type to maximize memory efficiency
-// goal: generate ld.128 and st.128 instructions
-template <typename T>
-struct packed_t {
-  // the (P)acked type for load/store
-  using P = array_t<T, 16 / sizeof(T)>;
-  // the (A)ccumulator type for reduction
-  using A = array_t<float, 16 / sizeof(T)>;
-};
-
-#define DINLINE __device__ __forceinline__
-
-// scalar cast functions
-DINLINE float upcast_s(half val) { return __half2float(val); }
-
-template <typename T>
-DINLINE T downcast_s(float val);
-template <>
-DINLINE half downcast_s(float val) {
-  return __float2half(val);
-}
-
-// scalar add functions
-// for some reason when compiling with Pytorch, the + operator for half and
-// bfloat is disabled so we call the intrinsics directly
-DINLINE half& assign_add(half& a, half b) {
-  a = __hadd(a, b);
-  return a;
-}
-DINLINE float& assign_add(float& a, float b) { return a += b; }
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
-template <>
-DINLINE nv_bfloat16 downcast_s(float val) {
-  return __float2bfloat16(val);
-}
-DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
-  a = __hadd(a, b);
-  return a;
-}
-#endif
-
-template <typename T, int N>
-DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
-#pragma unroll
-  for (int i = 0; i < N; i++) {
-    assign_add(a.data[i], b.data[i]);
-  }
-  return a;
-}
-
-template <typename T, int N>
-DINLINE array_t<float, N> upcast(array_t<T, N> val) {
-  if constexpr (std::is_same<T, float>::value) {
-    return val;
-  } else {
-    array_t<float, N> out;
-#pragma unroll
-    for (int i = 0; i < N; i++) {
-      out.data[i] = upcast_s(val.data[i]);
-    }
-    return out;
-  }
-}
-
-template <typename O>
-DINLINE O downcast(array_t<float, O::size> val) {
-  if constexpr (std::is_same<typename O::type, float>::value) {
-    return val;
-  } else {
-    O out;
-#pragma unroll
-    for (int i = 0; i < O::size; i++) {
-      out.data[i] = downcast_s<typename O::type>(val.data[i]);
-    }
-    return out;
-  }
-}
-
-#if !defined(USE_ROCM)
-
-static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
-               "l"(flag_addr));
-  #else
-  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
-               "l"(flag_addr));
-  #endif
-}
-
-static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
-  FlagType flag;
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  #else
-  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  #endif
-  return flag;
-}
-
-static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
-  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
-}
-
-static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
-  FlagType flag;
-  asm volatile("ld.volatile.global.u32 %0, [%1];"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  return flag;
-}
-
-// This function is meant to be used as the first synchronization in the all
-// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
-// prior memory accesses. Note: volatile writes will not be reordered against
-// other volatile writes.
-template <int ngpus>
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
-                              int rank) {
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
-    // Write the expected counter value to peer and wait for correct value
-    // from peer.
-    st_flag_volatile(peer_counter_ptr, flag);
-    while (ld_flag_volatile(self_counter_ptr) != flag);
-  }
-  __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus>
-DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
-                                      int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
-    st_flag_release(peer_counter_ptr, flag);
-    while (ld_flag_acquire(self_counter_ptr) != flag);
-  }
-  __syncthreads();
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-// This function is meant to be used as the second or the final
-// synchronization barrier in the all reduce kernel. If it's the final
-// synchronization barrier, we don't need to make any visibility guarantees
-// for prior memory accesses.
-template <int ngpus, bool final_sync = false>
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
-    // Write the expected counter value to peer and wait for correct value from
-    // peer.
-    if constexpr (!final_sync) {
-      st_flag_release(peer_counter_ptr, flag);
-      while (ld_flag_acquire(self_counter_ptr) != flag);
-    } else {
-      st_flag_volatile(peer_counter_ptr, flag);
-      while (ld_flag_volatile(self_counter_ptr) != flag);
-    }
-  }
-  if constexpr (!final_sync) __syncthreads();
-
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-#else
-
-template <int ngpus>
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
-                              int rank) {
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    // simultaneously write to the corresponding flag of all ranks.
-    // Latency = 1 p2p write
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
-                            flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
-    // wait until we got true from all ranks
-    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
-                                  __ATOMIC_RELAXED,
-                                  __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus>
-DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
-                                      int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
-                            flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM);
-    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
-                                  __ATOMIC_ACQUIRE,
-                                  __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  __syncthreads();
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus, bool final_sync = false>
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    // simultaneously write to the corresponding flag of all ranks.
-    // Latency = 1 p2p write
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
-                            flag,
-                            final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
-                            __MEMORY_SCOPE_SYSTEM);
-    // wait until we got true from all ranks
-    while (
-        __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
-                               final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
-                               __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  if constexpr (!final_sync) __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-#endif
-
-template <typename P, int ngpus, typename A>
-DINLINE P packed_reduce(const P* ptrs[], int idx) {
-  A tmp = upcast(ptrs[0][idx]);
-#pragma unroll
-  for (int i = 1; i < ngpus; i++) {
-    packed_assign_add(tmp, upcast(ptrs[i][idx]));
-  }
-  return downcast<P>(tmp);
-}
-
-}  // namespace vllm
@@ -1,17 +0,0 @@
-#include "core/registration.h"
-#include "flash_kda.h"
-
-TORCH_LIBRARY(_flashkda_C, m) {
-  m.def("get_workspace_size(int T_total, int H, int N=1) -> int",
-        &get_workspace_size);
-  m.def(
-      "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
-      "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, "
-      "float lower_bound, "
-      "Tensor? initial_state=None, Tensor(b!)? final_state=None, "
-      "Tensor? cu_seqlens=None) -> ()");
-}
-
-TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); }
-
-REGISTER_EXTENSION(_flashkda_C)
@@ -3,155 +3,19 @@

 #include <Python.h>

-#include <errno.h>
-#include <fcntl.h>
 #include <unistd.h>

-#include <filesystem>
-#include <string>
 #include <vector>

-#if defined(O_DIRECT)
-constexpr int kODirectFlag = O_DIRECT;
-#else
-constexpr int kODirectFlag = 0;
-#endif
-
 extern "C" {

-namespace {
-
-// Returns 0 on success, or the std::error_code's POSIX-compatible value on
-// failure, mirroring the errno convention used by the syscalls below.
-inline int ensure_parent_dirs(const std::string& path) {
-  const auto parent = std::filesystem::path(path).parent_path();
-  if (parent.empty()) {
-    return 0;
-  }
-  std::error_code ec;
-  std::filesystem::create_directories(parent, ec);
-  return ec ? ec.value() : 0;
-}
-
-// Core single-block store: src/size are raw pointer + byte count. Returns 0
-// on success, or the errno of the failing step on failure -- captured
-// before any subsequent cleanup call can overwrite it. On failure, the temp
-// file is removed.
-inline int _store_block(const char* tmp_path, const char* dest_path,
-                        const char* src, size_t size, bool use_o_direct) {
-  if (access(dest_path, F_OK) == 0) {
-    return 0;  // Already present.
-  }
-
-  if (const int err = ensure_parent_dirs(dest_path); err != 0) {
-    return err;
-  }
-
-  const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
-  const int fd = open(
-      tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644);
-  if (fd < 0) {
-    return errno;
-  }
-
-  const ssize_t written = write(fd, src, size);
-  if (written < 0 || static_cast<size_t>(written) != size) {
-    const int err = written < 0 ? errno : EIO;
-    close(fd);  // Best-effort cleanup; the real error is already captured.
-    unlink(tmp_path);
-    return err;
-  }
-
-  if (close(fd) != 0) {
-    const int err = errno;
-    unlink(tmp_path);
-    return err;
-  }
-
-  if (rename(tmp_path, dest_path) != 0) {
-    const int err = errno;
-    unlink(tmp_path);
-    return err;
-  }
-
-  return 0;
-}
-
-// Core single-block load: dst/size are raw pointer + byte count. Returns 0
-// on success, or the errno of the failing step on failure. On failure,
-// the source file is removed since a partially-read block should not be reused.
-inline int _load_block(const char* source_path, char* dst, size_t size,
-                       bool use_o_direct) {
-  const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
-  const int fd = open(source_path, O_RDONLY | o_direct_flag, 0);
-  if (fd < 0) {
-    const int err = errno;
-    unlink(source_path);
-    return err;
-  }
-
-  const ssize_t bytes_read = read(fd, dst, size);
-  if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) {
-    const int err = bytes_read < 0 ? errno : EIO;
-    close(fd);
-    unlink(source_path);
-    return err;
-  }
-
-  if (close(fd) != 0) {
-    const int err = errno;
-    unlink(source_path);
-    return err;
-  }
-
-  return 0;
-}
-
-inline void _batch_lookup(const std::vector<const char*>& paths,
+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;
  }
 }

-// Helper: extract a list[str] of length n into a vector<const char*>.
-// Returns false and sets a Python exception on error.
-inline bool extract_str_list(PyObject* list, Py_ssize_t n,
-                             std::vector<const char*>& out) {
-  for (Py_ssize_t i = 0; i < n; i++) {
-    out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr);
-    if (out[i] == nullptr) {
-      return false;
-    }
-  }
-  return true;
-}
-
-// Helper: extract a Py_buffer per element of a list[bytes-like] of length n.
-// On success, `out` holds n acquired buffers (caller must PyBuffer_Release
-// each). On failure, any buffers already acquired are released before
-// returning false, and a Python exception is set.
-inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags,
-                                std::vector<Py_buffer>& out) {
-  for (Py_ssize_t i = 0; i < n; i++) {
-    if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) {
-      for (Py_ssize_t j = 0; j < i; j++) {
-        PyBuffer_Release(&out[j]);
-      }
-      return false;
-    }
-  }
-  return true;
-}
-
-inline void release_buffer_list(std::vector<Py_buffer>& buffers) {
-  for (auto& buf : buffers) {
-    PyBuffer_Release(&buf);
-  }
-}
-
-}  // namespace
-
 /// @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.
@@ -187,157 +51,11 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
  return result;
 }

-/// @brief Store a batch of blocks, each from its own buffer, to disk.
-/// @param tmp_paths    list[str] – one temp path per block.
-/// @param dest_paths   list[str] – one destination path per block.
-/// @param buffers      list[bytes-like] – one source buffer per block.
-/// @param use_o_direct bool – whether to open files with O_DIRECT
-///                     (default True). Ignored where O_DIRECT is unsupported
-///                     by the platform.
-/// @note Releases the GIL for the entire batch. Raises on first error.
-static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) {
-  PyObject* tmp_paths_obj = nullptr;
-  PyObject* dest_paths_obj = nullptr;
-  PyObject* buffers_obj = nullptr;
-  int use_o_direct = 1;
-
-  if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj,
-                        &PyList_Type, &dest_paths_obj, &PyList_Type,
-                        &buffers_obj, &use_o_direct)) {
-    return nullptr;
-  }
-
-  const Py_ssize_t n = PyList_Size(tmp_paths_obj);
-  if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) {
-    PyErr_SetString(
-        PyExc_ValueError,
-        "tmp_paths, dest_paths and buffers must have the same length");
-    return nullptr;
-  }
-
-  std::vector<const char*> tmp_paths(n);
-  std::vector<const char*> dest_paths(n);
-
-  if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr;
-  if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr;
-
-  std::vector<Py_buffer> buffers(n);
-  if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) {
-    return nullptr;
-  }
-
-  Py_ssize_t failed_index = -1;
-  int failure_errno = 0;
-
-  {
-    Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
-      const char* buf = static_cast<const char*>(buffers[i].buf);
-      const int err =
-          _store_block(tmp_paths[i], dest_paths[i], buf,
-                       static_cast<size_t>(buffers[i].len), use_o_direct);
-      if (err != 0) {
-        failed_index = i;
-        failure_errno = err;
-        break;
-      }
-    }
-    Py_END_ALLOW_THREADS
-  }
-
-  release_buffer_list(buffers);
-
-  if (failed_index >= 0) {
-    // PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
-    errno = failure_errno;
-    return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
-                                          dest_paths[failed_index]);
-  }
-
-  Py_RETURN_NONE;
-}
-
-/// @brief Load a batch of blocks from disk, each into its own buffer.
-/// @param source_paths list[str] – one source path per block.
-/// @param buffers      list[writable bytes-like] – one destination buffer
-///                     per block.
-/// @param use_o_direct bool – whether to open files with O_DIRECT
-///                     (default True). Ignored where O_DIRECT is unsupported
-///                     by the platform.
-/// @note Releases the GIL for the entire batch. Raises on first error.
-static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) {
-  PyObject* source_paths_obj = nullptr;
-  PyObject* buffers_obj = nullptr;
-  int use_o_direct = 1;
-
-  if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj,
-                        &PyList_Type, &buffers_obj, &use_o_direct)) {
-    return nullptr;
-  }
-
-  const Py_ssize_t n = PyList_Size(source_paths_obj);
-  if (PyList_Size(buffers_obj) != n) {
-    PyErr_SetString(PyExc_ValueError,
-                    "source_paths and buffers must have the same length");
-    return nullptr;
-  }
-
-  std::vector<const char*> source_paths(n);
-  if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr;
-
-  std::vector<Py_buffer> buffers(n);
-  if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) {
-    return nullptr;
-  }
-
-  Py_ssize_t failed_index = -1;
-  int failure_errno = 0;
-
-  {
-    Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
-      char* buf = static_cast<char*>(buffers[i].buf);
-      const int err =
-          _load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len),
-                      use_o_direct);
-      if (err != 0) {
-        failed_index = i;
-        failure_errno = err;
-        break;
-      }
-    }
-    Py_END_ALLOW_THREADS
-  }
-
-  release_buffer_list(buffers);
-
-  if (failed_index >= 0) {
-    // PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
-    errno = failure_errno;
-    return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
-                                          source_paths[failed_index]);
-  }
-
-  Py_RETURN_NONE;
-}
-
 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."},
-    {"batch_store_block", batch_store_block, METH_VARARGS,
-     "batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n"
-     "                  buffers: list[bytes-like],\n"
-     "                  use_o_direct: bool = True) -> None\n"
-     "\n"
-     "Store a batch of blocks, each from its own buffer, to disk. Raises on "
-     "first error."},
-    {"batch_load_block", batch_load_block, METH_VARARGS,
-     "batch_load_block(source_paths: list[str],\n"
-     "                 buffers: list[writable bytes-like],\n"
-     "                 use_o_direct: bool = True) -> None\n"
-     "\n"
-     "Load a batch of blocks from disk into corresponding buffers. "
-     "Raises on first error."},
    {nullptr, nullptr, 0, nullptr},
 };

@@ -464,66 +464,6 @@ __global__ void swigluoai_and_mul_kernel(
  }
 }

-// SITU (Kimi SituGLU) gated activation. Non-interleaved layout:
-// input = [gate(d), up(d)] per token.
-//   gate_out = beta * tanh(gate / beta) * sigmoid(gate)
-//   up_out   = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up
-//   out      = gate_out * up_out
-// Compute is done in fp32 and written straight to `out` -- no intermediate
-// tensors and no full-tensor fp32 upcast (the pure-torch forward_native
-// allocated ~8 fp32 temporaries per call, which blows up MoE profiling).
-template <typename scalar_t>
-__global__ void situ_and_mul_kernel(
-    scalar_t* __restrict__ out,          // [..., d]
-    const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float beta, const float linear_beta) {
-  const int64_t row = blockIdx.x;
-  const scalar_t* gate_ptr = input + row * 2 * d;
-  const scalar_t* up_ptr = gate_ptr + d;
-  scalar_t* out_ptr = out + row * d;
-  const bool clamp_up = linear_beta > 0.0f;
-  const float inv_beta = 1.0f / beta;
-  const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
-  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
-    const float g = (float)VLLM_LDG(&gate_ptr[idx]);
-    const float u = (float)VLLM_LDG(&up_ptr[idx]);
-    const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
-    const float up_out =
-        clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
-    out_ptr[idx] = (scalar_t)(gate_out * up_out);
-  }
-}
-
-template <typename scalar_t>
-__global__ void masked_situ_and_mul_kernel(
-    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
-    const int* __restrict__ expert_num_tokens, const int max_num_tokens,
-    const int d, const float beta, const float linear_beta) {
-  const int expert = blockIdx.y;
-  const int num_tokens = expert_num_tokens[expert];
-  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
-  if (idx >= d || num_tokens == 0) {
-    return;
-  }
-
-  const bool clamp_up = linear_beta > 0.0f;
-  const float inv_beta = 1.0f / beta;
-  const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
-  const int64_t expert_row = static_cast<int64_t>(expert) * max_num_tokens;
-  for (int token = 0; token < num_tokens; ++token) {
-    const int64_t row = expert_row + token;
-    const scalar_t* gate_ptr = input + row * 2 * d;
-    const scalar_t* up_ptr = gate_ptr + d;
-    scalar_t* out_ptr = out + row * d;
-    const float g = (float)VLLM_LDG(&gate_ptr[idx]);
-    const float u = (float)VLLM_LDG(&up_ptr[idx]);
-    const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
-    const float up_out =
-        clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
-    out_ptr[idx] = (scalar_t)(gate_out * up_out);
-  }
-}
-
 }  // namespace vllm

 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -613,54 +553,6 @@ void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
                       double alpha, double limit) {
  LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
 }
-
-// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed
-// through), matching SituAndMul(linear_beta=None) on the Python side.
-void situ_and_mul(torch::stable::Tensor& out,    // [..., d]
-                  torch::stable::Tensor& input,  // [..., 2 * d]
-                  double beta, double linear_beta) {
-  int d = input.size(-1) / 2;
-  int64_t num_tokens = input.numel() / input.size(-1);
-  if (num_tokens == 0) {
-    return;
-  }
-  dim3 grid(num_tokens);
-  dim3 block(std::min(d, 1024));
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-      input.scalar_type(), "situ_and_mul_kernel", [&] {
-        vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
-            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
-            d, (float)beta, (float)linear_beta);
-      });
-}
-
-void masked_situ_and_mul(torch::stable::Tensor& out,    // [E, T, d]
-                         torch::stable::Tensor& input,  // [E, T, 2 * d]
-                         const torch::stable::Tensor& expert_num_tokens,
-                         double beta, double linear_beta) {
-  int num_experts = input.size(0);
-  int max_num_tokens = input.size(1);
-  int d = input.size(2) / 2;
-  if (num_experts == 0 || max_num_tokens == 0) {
-    return;
-  }
-  constexpr int block_size = 256;
-  dim3 grid((d + block_size - 1) / block_size, num_experts);
-  dim3 block(block_size);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-      input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
-        vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
-            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
-            expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
-            (float)beta, (float)linear_beta);
-      });
-}
 namespace vllm {

 // Element-wise activation kernel template.
@@ -21,10 +21,7 @@ __global__ void merge_attn_states_kernel(
    const float* prefix_lse, const scalar_t* suffix_output,
    const float* suffix_lse, const uint num_tokens, const uint num_heads,
    const uint head_size, const uint prefix_head_stride,
-    const uint output_head_stride, const uint prefix_lse_head_stride,
-    const uint prefix_lse_token_stride, const uint suffix_lse_head_stride,
-    const uint suffix_lse_token_stride, const uint output_lse_head_stride,
-    const uint output_lse_token_stride, const uint prefix_num_tokens,
+    const uint output_head_stride, const uint prefix_num_tokens,
    const float* output_scale) {
  // Inputs always load 128-bit packs (pack_size elements of scalar_t).
  // Outputs store pack_size elements of output_t, which is smaller for FP8.
@@ -87,19 +84,15 @@ __global__ void merge_attn_states_kernel(
      }
    }
    if (output_lse != nullptr && pack_idx == 0) {
-      float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
-                               token_idx * suffix_lse_token_stride];
-      output_lse[head_idx * output_lse_head_stride +
-                 token_idx * output_lse_token_stride] = s_lse;
+      float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
+      output_lse[head_idx * num_tokens + token_idx] = s_lse;
    }
    return;
  }

  // For tokens within prefix range, merge prefix and suffix
-  float p_lse = prefix_lse[head_idx * prefix_lse_head_stride +
-                           token_idx * prefix_lse_token_stride];
-  float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
-                           token_idx * suffix_lse_token_stride];
+  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
+  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
  p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
  s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;

@@ -139,8 +132,7 @@ __global__ void merge_attn_states_kernel(
    }
    // We only need to write to output_lse once per head.
    if (output_lse != nullptr && pack_idx == 0) {
-      output_lse[head_idx * output_lse_head_stride +
-                 token_idx * output_lse_token_stride] = max_lse;
+      output_lse[head_idx * num_tokens + token_idx] = max_lse;
    }
    return;
  }
@@ -195,8 +187,7 @@ __global__ void merge_attn_states_kernel(
  // We only need to write to output_lse once per head.
  if (output_lse != nullptr && pack_idx == 0) {
    float out_lse = logf(out_se) + max_lse;
-    output_lse[head_idx * output_lse_head_stride +
-               token_idx * output_lse_token_stride] = out_lse;
+    output_lse[head_idx * num_tokens + token_idx] = out_lse;
  }
 }

@@ -230,9 +221,6 @@ __global__ void merge_attn_states_kernel(
            reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
            reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
            num_heads, head_size, prefix_head_stride, output_head_stride,   \
-            prefix_lse_head_stride, prefix_lse_token_stride,                \
-            suffix_lse_head_stride, suffix_lse_token_stride,                \
-            output_lse_head_stride, output_lse_token_stride,                \
            prefix_num_tokens, output_scale_ptr);                           \
  }

@@ -271,19 +259,6 @@ void merge_attn_states_launcher(
  const uint head_size = output.size(2);
  const uint prefix_head_stride = prefix_output.stride(1);
  const uint output_head_stride = output.stride(1);
-  // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views
-  // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index
-  // them by their actual strides rather than assuming a contiguous layout.
-  const uint prefix_lse_head_stride = prefix_lse.stride(0);
-  const uint prefix_lse_token_stride = prefix_lse.stride(1);
-  const uint suffix_lse_head_stride = suffix_lse.stride(0);
-  const uint suffix_lse_token_stride = suffix_lse.stride(1);
-  uint output_lse_head_stride = 0;
-  uint output_lse_token_stride = 0;
-  if (output_lse.has_value()) {
-    output_lse_head_stride = output_lse.value().stride(0);
-    output_lse_token_stride = output_lse.value().stride(1);
-  }
  // Thread mapping is based on input BF16 pack_size
  const uint pack_size = 16 / sizeof(scalar_t);
  STD_TORCH_CHECK(head_size % pack_size == 0,
--- a/Show More
+++ b/Show More