[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262)

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>

.buildkite/check-torch-abi.py

@@ -0,0 +1,103 @@
+# 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())

.buildkite/ci_config.yaml

@@ -14,6 +14,7 @@ run_all_patterns:
   - "setup.py"
   - "csrc/"
   - "cmake/"
+  - ".buildkite/check-torch-abi.py"
 run_all_exclude_patterns:
   - "docker/Dockerfile."
   - "csrc/cpu/"

.buildkite/hardware_tests/cpu.yaml

@@ -18,6 +18,8 @@ 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 "
@@ -28,7 +30,9 @@ 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/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"
 
 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests

.buildkite/intel_jobs/benchmarks_intel.yaml

@@ -0,0 +1,26 @@
+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/'

.buildkite/intel_jobs/engine_intel.yaml

@@ -2,6 +2,44 @@ 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
@@ -23,3 +61,41 @@ 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"'

.buildkite/intel_jobs/misc_intel.yaml

@@ -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 (2 GPUs)
+- label: NixlConnector PD accuracy (4 GPUs)
   timeout_in_minutes: 60
-  num_devices: 2
+  num_devices: 4
   device: intel_gpu
   agent_tags:
     label: production
-    gpu: 2+
+    gpu: 4+
     mem: 16+
   no_plugin: true
   working_dir: "."
@@ -148,7 +148,10 @@ steps:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
-      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
+      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'
 
 - label: Regression
   key: regression
@@ -259,3 +262,25 @@ 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'

.buildkite/intel_jobs/model_executor_intel.yaml

@@ -0,0 +1,33 @@
+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"'

.buildkite/intel_jobs/model_runner_v2_intel.yaml

@@ -8,7 +8,7 @@ steps:
   agent_tags:
     label: production
     gpu: 2+
-    mem: 16+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -28,7 +28,9 @@ 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)
@@ -60,3 +62,55 @@ 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'

.buildkite/intel_jobs/samplers_intel.yaml

@@ -0,0 +1,29 @@
+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'

.buildkite/intel_jobs/test-intel.yaml

@@ -145,7 +145,32 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        pytest -v -s quantization/test_auto_round.py'
+        pytest -v -s quantization/test_auto_round.py &&
+        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'
   - label: "XPU compressed tensors FP8 test"
     depends_on:
       - image-build-xpu
@@ -168,4 +193,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'

.buildkite/performance-benchmarks/tests/serving-tests.json

@@ -28,11 +28,6 @@
         "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": {

.buildkite/release-pipeline.yaml

@@ -813,8 +813,8 @@ steps:
 
             # 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" .
+            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

.buildkite/scripts/build-macos-wheel.sh

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

.buildkite/scripts/ci-bake-rocm.sh

@@ -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_rixl 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_nixl 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.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.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.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.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
+    metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_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 RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
+    for var in NIXL_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 rixl_branch=""
+    local nixl_branch=""
     local ucx_branch=""
     local rocshmem_branch=""
     local deepep_branch=""
-    local rixl_material=""
+    local nixl_material=""
     local rocshmem_material=""
     local deepep_material=""
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
     dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
-    rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
+    nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_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 "${rixl_branch}" && -n "${ucx_branch}" ]]; then
-        rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
-        RIXL_CACHE_KEY=$(
+    if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
+        nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
+        NIXL_CACHE_KEY=$(
             compose_dependency_cache_key \
-                "${rixl_branch}-ucx-${ucx_branch}" \
-                "${rixl_material}"
+                "${nixl_branch}-ucx-${ucx_branch}" \
+                "${nixl_material}"
         )
-        export RIXL_CACHE_KEY
-        echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
+        export NIXL_CACHE_KEY
+        echo "NIXL dependency cache key: ${NIXL_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
-        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:-}"
+        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:-}"
             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 rixl_ref=""
+    local nixl_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 rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
+            for target in nixl-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 "${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}"
+    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}"
         else
-            echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
-            add_dependency_cache_target "rixl-rocm-ci"
+            echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
+            add_dependency_cache_target "nixl-rocm-ci"
         fi
     fi
 

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

@@ -35,7 +35,7 @@ set -o pipefail
 : "${PY_COLORS:=1}"
 : "${ROCM_DOCKER_TTY:=1}"
 : "${PYTHONFAULTHANDLER:=1}"
-: "${PYTEST_TIMEOUT:=2100}"
+: "${PYTEST_TIMEOUT:=2400}"
 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 15 minutes, then stop an individual test after 35 minutes.
+# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
 if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
+  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
 fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
   " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
@@ -387,6 +387,7 @@ 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
@@ -400,16 +401,19 @@ 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_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
+  export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
   export PYTORCH_ROCM_ARCH=""
 
   mkdir -p "${TMPDIR}" \
@@ -417,7 +421,10 @@ initialize_native_environment() {
     "${TRITON_CACHE_DIR}" \
     "${VLLM_CACHE_ROOT}" \
     "${XDG_CACHE_HOME}" \
-    "${HF_HOME}" || return 1
+    "${HF_HOME}" \
+    "${HF_DATASETS_CACHE}" || return 1
+
+  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
 
   if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
     if ! command -v findmnt >/dev/null 2>&1; then
@@ -430,6 +437,18 @@ 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() {

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

@@ -1,10 +1,11 @@
 #!/bin/bash
 set -euox pipefail
 
-export VLLM_CPU_KVCACHE_SPACE=1 
+export VLLM_CPU_KVCACHE_SPACE=1
 export VLLM_CPU_CI_ENV=1
-# Reduce sub-processes for acceleration
-export TORCH_COMPILE_DISABLE=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.
 export VLLM_ENABLE_V1_MULTIPROCESSING=0
 
 SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -49,15 +50,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 > test_0.log 2>&1 &
+./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > 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 > test_1.log 2>&1 &
+./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > 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 > test_2.log 2>&1 &
+./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
 PID_TEST_2=$!
 
 wait_for_pid_and_check_log $PID_TEST_0 test_0.log

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

@@ -40,7 +40,9 @@ 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_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"
 
   # skip tests requiring model downloads if HF_TOKEN is not set
   # due to rate-limits
@@ -97,3 +99,4 @@ 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
+

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

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

.buildkite/test_areas/attention.yaml

@@ -16,8 +16,9 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 95
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 125
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/basic_correctness.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
   key: basic-correctness
-  timeout_in_minutes: 45
+  timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -18,7 +18,8 @@ steps:
   - pytest -v -s basic_correctness/test_cpu_offload.py
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd

.buildkite/test_areas/benchmarks.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
   key: benchmarks-cli-test
-  timeout_in_minutes: 30
+  timeout_in_minutes: 45
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -15,6 +15,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/cuda.yaml

@@ -16,6 +16,7 @@ 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
@@ -26,7 +27,10 @@ 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

.buildkite/test_areas/disaggregated.yaml

@@ -17,7 +17,7 @@ steps:
     amd:
       dind: false
       device: mi300_4
-      timeout_in_minutes: 85
+      timeout_in_minutes: 60
       depends_on:
         - image-build-amd
       source_file_dependencies:
@@ -68,7 +68,7 @@ steps:
     amd:
       dind: false
       device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 40
       depends_on:
         - image-build-amd
       source_file_dependencies:
@@ -94,7 +94,7 @@ steps:
     amd:
       dind: false
       device: mi300_4
-      timeout_in_minutes: 85
+      timeout_in_minutes: 60
       depends_on:
         - image-build-amd
       source_file_dependencies:
@@ -120,7 +120,7 @@ steps:
     amd:
       dind: false
       device: mi300_4
-      timeout_in_minutes: 80
+      timeout_in_minutes: 55
       depends_on:
         - image-build-amd
       source_file_dependencies:
@@ -131,6 +131,22 @@ 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
@@ -177,7 +193,7 @@ steps:
     amd:
       dind: false
       device: mi300_2
-      timeout_in_minutes: 70
+      timeout_in_minutes: 45
       depends_on:
         - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/distributed.yaml

@@ -41,6 +41,7 @@ steps:
     amd:
       dind: false
       device: mi300_2
+      timeout_in_minutes: 45
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/engine.yaml

@@ -28,8 +28,9 @@ steps:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 50
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
 
@@ -39,19 +40,21 @@ 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: mi325_1
-      timeout_in_minutes: 55
+      device: mi250_1
+      timeout_in_minutes: 45
       depends_on:
       - image-build-amd
 
 - label: e2e Scheduling (1 GPU)
   key: e2e-scheduling-1-gpu
-  timeout_in_minutes: 35
+  timeout_in_minutes: 53
   device: h200_18gb
   source_file_dependencies:
     - vllm/v1/
@@ -60,8 +63,8 @@ steps:
     - pytest -v -s v1/e2e/general/test_async_scheduling.py
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      device: mi250_1
+      timeout_in_minutes: 55
       depends_on:
       - image-build-amd
 
@@ -76,8 +79,8 @@ steps:
     - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 60
+      device: mi250_1
+      timeout_in_minutes: 50
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -116,6 +119,7 @@ steps:
     amd:
       dind: false
       device: mi300_2
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/entrypoints.yaml

@@ -30,16 +30,16 @@ steps:
   - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
   mirror:
     amd:
-      device: mi325_1
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
       depends_on:
       - image-build-amd
 
 - label: Entrypoints Integration (API Server)
   key: entrypoints-integration-api-server
   device: h200_35gb
-  timeout_in_minutes: 50
+  timeout_in_minutes: 75
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -52,14 +52,16 @@ steps:
   - pytest -v -s entrypoints/scale_out
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
       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: 45
+  timeout_in_minutes: 68
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -70,7 +72,8 @@ 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:
-      device: mi325_1
+      dind: false
+      device: mi300_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -78,7 +81,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 2)
   device: h200_35gb
   key: entrypoints-integration-api-server-openai-part-2
-  timeout_in_minutes: 45
+  timeout_in_minutes: 83
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -90,8 +93,9 @@ steps:
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 80
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 70
       depends_on:
       - image-build-amd
 
@@ -113,7 +117,8 @@ steps:
   - pytest -v -s entrypoints/anthropic
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -156,7 +161,7 @@ steps:
 - label: Entrypoints Integration (Pooling)
   device: h200_35gb
   key: entrypoints-integration-pooling
-  timeout_in_minutes: 50
+  timeout_in_minutes: 75
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -176,7 +181,9 @@ steps:
   - pytest -s entrypoints/openai/correctness/
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/expert_parallelism.yaml

@@ -18,6 +18,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -51,4 +52,5 @@ steps:
   - vllm/compilation/
   - tests/distributed/
   commands:
+  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
   - pytest -v -s distributed/test_elastic_ep.py

.buildkite/test_areas/fault_tolerance.yaml

@@ -0,0 +1,26 @@
+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

.buildkite/test_areas/kernels.yaml

@@ -61,9 +61,45 @@ 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
@@ -80,7 +116,8 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
       timeout_in_minutes: 90
       depends_on:
       - image-build-amd
@@ -118,7 +155,9 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 120
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization
@@ -148,8 +187,9 @@ steps:
   parallelism: 5
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 65
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
       source_file_dependencies:
       - csrc/quantization/cutlass_w8a8/moe/
       - csrc/moe/
@@ -174,17 +214,6 @@ 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

.buildkite/test_areas/lm_eval.yaml

@@ -14,8 +14,9 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 55
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 45
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -141,7 +142,7 @@ steps:
     amd:
       dind: false
       device: mi300_8
-      timeout_in_minutes: 60
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
       commands:
@@ -336,7 +337,7 @@ steps:
 
 - label: LM Eval KV-Offload (2xH100)
   key: kv-offload-medium
-  timeout_in_minutes: 30
+  timeout_in_minutes: 45
   device: h100
   num_devices: 2
   source_file_dependencies:
@@ -346,7 +347,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"
+  - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
 
 - label: LM Eval KV-Offload (4xH100)
   key: kv-offload-large

.buildkite/test_areas/lora.yaml

@@ -14,9 +14,10 @@ steps:
   parallelism: 4
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
       working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 65
+      timeout_in_minutes: 85
       source_file_dependencies:
       - vllm/lora
       - tests/lora
@@ -46,4 +47,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

.buildkite/test_areas/misc.yaml

@@ -25,13 +25,13 @@ steps:
     amd:
       dind: false
       device: mi300_1
-      timeout_in_minutes: 75
+      timeout_in_minutes: 50
       depends_on:
       - image-build-amd
 
 - label: V1 Sample + Logits
   key: v1-sample-logits
-  timeout_in_minutes: 45
+  timeout_in_minutes: 83
   device: h200_18gb
   source_file_dependencies:
     - vllm/config/
@@ -59,7 +59,9 @@ steps:
     - pytest -v -s v1/test_outputs.py
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 70
       depends_on:
       - image-build-amd
 
@@ -90,6 +92,7 @@ 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
@@ -103,6 +106,7 @@ 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
@@ -111,8 +115,9 @@ steps:
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 75
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
       depends_on:
       - image-build-amd
 
@@ -143,6 +148,8 @@ 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
 
@@ -206,7 +213,7 @@ steps:
   - vllm/multimodal
   - examples/
   commands:
-    - pip install tensorizer # for tensorizer test
+    - pip install --no-deps tensorizer # for tensorizer test
      # for basic
     - python3 basic/offline_inference/chat.py
     - python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -230,7 +237,9 @@ 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:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 75
       source_file_dependencies:
       - vllm/entrypoints
       - vllm/multimodal
@@ -257,6 +266,7 @@ steps:
   - vllm/utils/
   - vllm/v1/
   - tests/v1/tracing
+  - tests/tracing/
   commands:
   - "pip install \
       'opentelemetry-sdk>=1.26.0' \
@@ -264,12 +274,14 @@ 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:
-      device: mi325_2
+      dind: false
+      device: mi300_2
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
-      optional: true
 
 - label: Python-only Installation
   key: python-only-installation
@@ -284,8 +296,9 @@ steps:
   - bash standalone_tests/python_only_compile.sh
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 45
+      device: mi250_1
+      timeout_in_minutes: 55
+      soft_fail: true
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -385,7 +398,7 @@ steps:
 
 - label: Batch Invariance (A100)
   key: batch-invariance-a100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   device: a100
   source_file_dependencies:
     - vllm/v1/attention
@@ -395,11 +408,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[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 -k TRITON_MLA
 
 - label: Batch Invariance (H100)
   key: batch-invariance-h100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   device: h100
   source_file_dependencies:
     - vllm/v1/attention
@@ -410,12 +423,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[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]
+    - 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
 
 - label: Batch Invariance (B200)
   key: batch-invariance-b200
-  timeout_in_minutes: 35
+  timeout_in_minutes: 45
   device: b200-k8s
   source_file_dependencies:
     - vllm/v1/attention
@@ -426,11 +439,14 @@ 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[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]
+    - 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
     - 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

.buildkite/test_areas/model_executor.yaml

@@ -5,12 +5,14 @@ steps:
 - label: Model Executor
   device: h200_35gb
   key: model-executor
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
   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
@@ -28,13 +30,16 @@ 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

.buildkite/test_areas/model_runner_v2.yaml

@@ -41,7 +41,7 @@ steps:
   commands:
     - set -x
     - export VLLM_USE_V2_MODEL_RUNNER=1
-    - pip install tensorizer # for tensorizer test
+    - pip install --no-deps 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

.buildkite/test_areas/models_basic.yaml

@@ -42,10 +42,39 @@ steps:
     - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 50
       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:
@@ -55,7 +84,8 @@ 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/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/

.buildkite/test_areas/models_language.yaml

@@ -17,6 +17,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 45
       depends_on:
       - image-build-amd
 
@@ -39,6 +40,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -61,7 +63,6 @@ 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.
@@ -69,8 +70,9 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
       commands:
@@ -102,7 +104,6 @@ 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)'
@@ -130,14 +131,15 @@ steps:
     - pytest -v -s models/language/pooling -m 'not core_model'
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 120
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 95
       depends_on:
       - image-build-amd
 
 - label: Language Models Test (MTEB)
   key: language-models-test-mteb
-  timeout_in_minutes: 45
+  timeout_in_minutes: 68
   device: h200_18gb
   optional: true
   source_file_dependencies:

.buildkite/test_areas/models_multimodal.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
   key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 45
+  timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -14,13 +14,15 @@ steps:
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
       depends_on:
       - image-build-amd
 
 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
   key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 50
+  timeout_in_minutes: 75
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -31,7 +33,9 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
       depends_on:
       - image-build-amd
 
@@ -47,14 +51,15 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
-      device: mi325_1
+      device: mi250_1
+      timeout_in_minutes: 55
       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: 50
+  timeout_in_minutes: 75
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
@@ -65,7 +70,9 @@ 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:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 50
       depends_on:
       - image-build-amd
 
@@ -85,7 +92,7 @@ steps:
 
 - label: Multi-Modal Processor # 44min
   key: multi-modal-processor
-  timeout_in_minutes: 65
+  timeout_in_minutes: 98
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -109,6 +116,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 35
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -131,7 +139,9 @@ steps:
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 90
       depends_on:
       - image-build-amd
 
@@ -166,8 +176,9 @@ steps:
     - pytest -v -s models/multimodal/pooling -m 'not core_model'
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 75
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/pytorch.yaml

@@ -5,7 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
   device: h200_35gb
   key: pytorch-compilation-unit-tests
-  timeout_in_minutes: 90
+  timeout_in_minutes: 150
   source_file_dependencies:
     - vllm/__init__.py
     - vllm/_aiter_ops.py
@@ -107,13 +107,6 @@ 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
@@ -236,6 +229,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/quantization.yaml

@@ -24,8 +24,7 @@ steps:
   - 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' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 8
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
 
 - label: Quantized Fusions
   device: h200_35gb
@@ -68,5 +67,4 @@ steps:
   - vllm/model_executor/layers/quantization
   - tests/models/quantization
   commands:
-    - pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 3
+    - pytest -v -s models/quantization

.buildkite/test_areas/rust_frontend.yaml

@@ -26,7 +26,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - 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_logit_bias_validation.py -k "not multiple"
 
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"

.buildkite/test_areas/samplers.yaml

@@ -19,8 +19,18 @@ steps:
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
   mirror:
     amd:
-      device: mi325_1
+      device: mi250_1
+      timeout_in_minutes: 40
       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

.buildkite/test_areas/spec_decode.yaml

@@ -14,8 +14,9 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 60
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -53,8 +54,9 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
   mirror:
     amd:
-      device: mi325_1
-      timeout_in_minutes: 65
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 75
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -88,14 +90,15 @@ 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:
-      device: mi325_1
-      timeout_in_minutes: 55
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 35
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -119,7 +122,8 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
       timeout_in_minutes: 55
       depends_on:
       - image-build-amd
@@ -170,3 +174,31 @@ 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"

.buildkite/test_areas/torch_abi.yaml

@@ -0,0 +1,14 @@
+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

.buildkite/test_areas/weight_loading.yaml

@@ -17,6 +17,7 @@ steps:
     amd:
       dind: false
       device: mi300_2
+      timeout_in_minutes: 35
       depends_on:
       - image-build-amd
       commands:

.github/CODEOWNERS

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

.github/mergify.yml

@@ -19,6 +19,7 @@ pull_request_rules:
   description: Comment on PR when pre-commit check fails
   conditions:
     - check-failure=pre-commit
+    - -check-cancelled=pre-commit
     - -closed
     - -draft
     - or:
@@ -181,6 +182,18 @@ 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:
@@ -220,6 +233,31 @@ 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:

.github/workflows/issue_autolabel.yml

@@ -130,6 +130,66 @@ 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: [...],
@@ -491,4 +551,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(', ')}`);

.github/workflows/new_pr_bot.yml

@@ -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. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
+                  '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.',
                   '',
-                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
+                  '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.',
                   '',
                   'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
                   '',

.github/workflows/pre-commit.yml

@@ -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 (the ready labels also trigger tests) 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 to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`);
           }
 
   pre-commit:

.github/workflows/run-ci-command.yml

@@ -0,0 +1,40 @@
+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 }}

.github/workflows/scripts/run_ci_command.py

@@ -0,0 +1,581 @@
+# 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()

.github/workflows/scripts/test_run_ci_command.py

@@ -0,0 +1,363 @@
+# 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()

.gitignore

@@ -173,9 +173,6 @@ venv.bak/
 
 # mkdocs documentation
 /site
-docs/argparse
-docs/examples/*
-!docs/examples/README.md
 
 # mypy
 .mypy_cache/

.markdownlint.yaml

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

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_install_hook_types:
 default_stages:
   - pre-commit # Run locally
   - manual # Run in CI
-exclude: 'vllm/third_party/.*'
+exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.14.0
@@ -260,10 +260,6 @@ 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

CMakeLists.txt

@@ -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.11.0")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.13.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,6 +114,11 @@ 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
@@ -214,10 +219,8 @@ 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. 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.
+  # `+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.
   #
   clear_cuda_arches(CUDA_ARCH_FLAGS)
   extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")
@@ -227,6 +230,13 @@ 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
@@ -411,8 +421,11 @@ 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_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
+    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA" AND
      DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -420,7 +433,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;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+        "9.0a;10.0f;10.1f;10.3f;10.7f;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}")
@@ -695,7 +708,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;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;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()
@@ -815,7 +828,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;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
   else()
     cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
@@ -899,7 +912,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;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
   else()
     cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
@@ -924,7 +937,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;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;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()
@@ -981,7 +994,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;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
   else()
     cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
@@ -1047,7 +1060,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;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
   else()
     cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
   endif()
@@ -1069,6 +1082,41 @@ 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)
@@ -1110,6 +1158,14 @@ 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)
@@ -1407,6 +1463,7 @@ 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)
 

README.md

@@ -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, 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, 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.
 
 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
 

benchmarks/attention_benchmarks/benchmark.py

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

benchmarks/kernels/benchmark_inkling_qkvr_prep.py

@@ -0,0 +1,176 @@
+# 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())

benchmarks/kernels/benchmark_k3_cutedsl_residual.py

@@ -0,0 +1,367 @@
+# 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()

benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py

@@ -0,0 +1,806 @@
+# 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()

benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py

@@ -0,0 +1,239 @@
+# 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()

benchmarks/kernels/cpu/benchmark_cpu_attn.py

@@ -154,7 +154,7 @@ def main(
                 scale=scale,
                 causal=True,
                 alibi_slopes=None,
-                sliding_window=window_size,
+                sliding_window=window_size if sliding_window is not None else -1,
                 block_table=block_tables,
                 softcap=0,
                 scheduler_metadata=metadata,

benchmarks/replayssm/e2e_decode_speedup.py

@@ -0,0 +1,267 @@
+# 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()

cmake/cpu_extension.cmake

@@ -15,6 +15,7 @@ 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")
@@ -96,12 +97,14 @@ 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
@@ -111,6 +114,11 @@ 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.
@@ -166,6 +174,11 @@ 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")
@@ -430,6 +443,7 @@ 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")
@@ -446,8 +460,13 @@ 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)	
@@ -489,6 +508,7 @@ 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"
@@ -502,6 +522,7 @@ 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

cmake/external_projects/deepgemm.cmake

@@ -68,6 +68,9 @@ 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()

cmake/external_projects/flashkda.cmake

@@ -0,0 +1,74 @@
+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()

cmake/external_projects/flashmla.cmake

@@ -60,6 +60,9 @@ 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()
@@ -188,4 +191,3 @@ else()
     add_custom_target(_flashmla_C)
     add_custom_target(_flashmla_extension_C)
 endif()
-

cmake/external_projects/fmha_sm100.cmake

@@ -17,7 +17,7 @@ else()
   FetchContent_Declare(
     fmha_sm100
     GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
+    GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""

cmake/external_projects/qutlass.cmake

@@ -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 "830d2c4537c7396e14a02a46fbddd18b5d107c65")
+  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
 
   set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _qutlass_fc_root)
@@ -55,7 +55,11 @@ 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}")
-  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.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()
 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}")
@@ -125,8 +129,6 @@ 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
@@ -139,9 +141,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
     WITH_SOABI)
 
   target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_DISABLE_PYBIND=1
+    QUTLASS_MINIMAL_BUILD=1
     TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
+    TORCH_TARGET_VERSION=0x020B000000000000ULL
+    USE_CUDA)
 
   set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
     $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>

cmake/external_projects/vllm_flash_attn.cmake

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

cmake/utils.cmake

@@ -241,14 +241,15 @@ endmacro()
 # `<major>.<minor>`, dedupes them and then sorts them in ascending order and 
 # stores them in `OUT_ARCHES`.
 #
-# 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"
+# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags.
+# This handles mismatches such as `arch=compute_20,code=sm_121`.
 function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
   set(_CUDA_ARCHES)
   foreach(_ARCH ${CUDA_ARCH_FLAGS})
-    string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
+    string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
+    if (NOT _COMPUTE)
+      string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
+    endif()
     if (_COMPUTE)
       set(_COMPUTE ${CMAKE_MATCH_1})
     endif()
@@ -396,14 +397,24 @@ 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}" 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)
+      if("${_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)
@@ -487,7 +498,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;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;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()

csrc/cpu/cpu_arch_macros.h

@@ -172,4 +172,15 @@
 
 #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

csrc/cpu/cpu_attn.cpp

@@ -30,7 +30,8 @@ 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::optional<torch::Tensor>& dynamic_causal,
+    const std::string& kv_cache_dtype) {
   cpu_attention::ISA isa;
   if (isa_hint == "amx") {
     isa = cpu_attention::ISA::AMX;
@@ -65,9 +66,11 @@ 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, 0, [&]() {
-      input.elem_size = sizeof(scalar_t);
+    CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
+      input.elem_size = sizeof(attn_impl::kv_cache_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 =

csrc/cpu/cpu_attn_vec.hpp

@@ -102,7 +102,9 @@ class TileGemm82 {
     kv_cache_t* __restrict__ curr_b = b_tile;
 
     for (int32_t k = 0; k < dynamic_k_size; ++k) {
-      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
+      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;
 
       float* __restrict__ curr_m_a = curr_a;
       vec_op::unroll_loop<int32_t, M>([&](int32_t i) {

csrc/cpu/cpu_fused_moe.cpp

@@ -1,5 +1,6 @@
 #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"
 
@@ -43,193 +44,9 @@
   }()
 
 namespace {
-enum class FusedMOEAct {
-  SiluAndMul,
-  SwigluOAIAndMul,
-  GeluAndMul,
-  GeluTanhAndMul,
-};
 
-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.");
-  }
-}
+using cpu_fused_moe_utils::apply_gated_act;
+using cpu_fused_moe_utils::FusedMOEAct;
 
 template <typename scalar_t, typename gemm_t>
 void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
     }
   }
 }
+
 }  // namespace
 
 void prepack_moe_weight(
@@ -864,7 +682,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 = get_act_type(act);
+  const FusedMOEAct act_type = cpu_fused_moe_utils::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");

csrc/cpu/cpu_fused_moe_activations.hpp

@@ -0,0 +1,204 @@
+// 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

csrc/cpu/cpu_fused_moe_int8.cpp

@@ -0,0 +1,647 @@
+// 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);
+        });
+      });
+}

csrc/cpu/cpu_types_vsx.hpp

@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
 
   explicit FP32Vec4(__vector float data) : reg(data) {}
 
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  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) {}
 
-  explicit FP32Vec8(const FP32Vec8& data) {
+  FP32Vec8(const FP32Vec8& data) {
     reg.val[0] = data.reg.val[0];
     reg.val[1] = data.reg.val[1];
   }
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = fp16_to_fp32_bits(raw_lo);
   }
   float reduce_sum() const {
-    AliasReg ar;
-    ar.reg = reg;
-    float result = 0;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, &ar](int i) { result += ar.values[i]; });
-
-    return result;
+    // 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);
   }
   FP32Vec8 exp() const {
     f32x4x2_t out;
@@ -592,7 +593,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
   explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
 
-  explicit FP32Vec16(const FP32Vec16& data) {
+  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];
@@ -746,6 +747,15 @@ 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]);

csrc/cpu/cpu_types_vxe.hpp

@@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
 
   explicit FP32Vec4(__vector float data) : reg(data) {}
 
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  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) {}
 
-  explicit FP32Vec8(const FP32Vec8& data) {
+  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) {}
 
-  explicit FP32Vec16(const FP32Vec16& data) {
+  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];

csrc/cpu/mamba_cpu.cpp

@@ -0,0 +1,285 @@
+// 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);
+  });
+}

csrc/cpu/mamba_kernels.hpp

@@ -0,0 +1,382 @@
+// 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

csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp

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

csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp

@@ -0,0 +1,424 @@
+// 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

csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp

@@ -1,3 +1,6 @@
+// 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
 
@@ -16,9 +19,6 @@ 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 += Nr * K;
+    b_tile += 4 * TileSize;
   }
 
   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 += Nr * K;
-    b_tile1 += Nr * K;
+    b_tile0 += 4 * TileSize;
+    b_tile1 += 4 * TileSize;
   }
 
   store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
@@ -223,6 +223,9 @@ 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;
@@ -246,6 +249,9 @@ 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;
@@ -253,7 +259,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
 
  public:
   // physical layout [
-  //  M / 8; Mr is 8
+  //  M / (8 or 4); Mr is 8 or 4
   //  K / 4; K for bfmmla is 4
   //  4,   ; 4 row-pairs for each 8 rows
   //  2,   ; row-pair is 2 rows
@@ -439,7 +445,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 += Nr_gemv) {
+    for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
       const bfloat16_t* __restrict__ b_panel =
           reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
 

csrc/cpu/sgl-kernels/conv.cpp

@@ -451,6 +451,90 @@ 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]
@@ -545,7 +629,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, width - 1) itype
+//   conv_states: (..., dim, state_len) itype, where state_len >= width - 1
 //   activation: either None or "silu" or "swish"
 //   pad_slot_id: int
 //
@@ -586,11 +670,14 @@ 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);
-    CHECK_EQ(conv_states_val.size(2), width - 1);
+    const int64_t state_len = conv_states_val.size(2);
+    CHECK_GE(state_len, 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);
@@ -651,14 +738,14 @@ at::Tensor causal_conv1d_fwd_cpu(
 
 // API aligned with GPUs
 //
-//   x: (batch, dim) or (batch, dim, seqlen)
+//   x: (batch, dim) or (batch, seqlen, dim)
 //   conv_state: (..., dim, state_len), where state_len >= width - 1
 //   weight: (dim, width)
 //   bias: (dim,)
-//   cache_seqlens: (batch,), dtype int32.
+//   num_accepted_tokens: (batch,), dtype int32.
 //   conv_state_indices: (batch,), dtype int32
 //   pad_slot_id: int
-//   out: (batch, dim) or (batch, dim, seqlen)
+//   out: (batch, dim) or (batch, seqlen, dim)
 //
 at::Tensor causal_conv1d_update_cpu(
     const at::Tensor& x,
@@ -666,7 +753,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>& cache_seqlens,
+    const std::optional<at::Tensor>& num_accepted_tokens,
     const std::optional<at::Tensor>& conv_state_indices,
     int64_t pad_slot_id,
     bool is_vnni) {
@@ -674,13 +761,13 @@ at::Tensor causal_conv1d_update_cpu(
   CHECK_CONTIGUOUS(weight);
   auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
 
-  // 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.");
+  TORCH_CHECK(
+      x.dim() == 2 || x.dim() == 3,
+      "causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");
 
   int64_t batch = x.size(0);
-  int64_t dim = x.size(1);
-  int64_t seqlen = 1;
+  int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);
+  int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);
   int64_t width = weight.size(-1);
 
   const auto scalar_type = x.scalar_type();
@@ -690,10 +777,84 @@ at::Tensor causal_conv1d_update_cpu(
 
   CHECK_EQ(conv_states.scalar_type(), scalar_type);
   CHECK_EQ(conv_states.size(1), dim);
-  CHECK_EQ(conv_states.size(2), width - 1);
+  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.");
 
   // 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});

csrc/cpu/torch_bindings.cpp

@@ -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>& cache_seqlens,
+    bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,
     const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
     bool is_vnni);
 
@@ -163,7 +163,8 @@ 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::optional<torch::Tensor>& dynamic_causal,
+    const std::string& kv_cache_dtype);
 
 void cpu_attn_reshape_and_cache(const torch::Tensor& key,
                                 const torch::Tensor& value,
@@ -207,12 +208,52 @@ 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 {
@@ -476,7 +517,8 @@ 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? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
+      "Tensor? num_accepted_tokens, 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
@@ -536,7 +578,8 @@ 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) -> Tensor",
+      "enable_kv_split, Tensor? dynamic_causal, "
+      "str kv_cache_dtype=\"auto\") -> Tensor",
       &get_scheduler_metadata);
   ops.def(
       "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -570,8 +613,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif
 
   // fused moe
-#if defined(__AVX512F__) || \
-    (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
+#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
   ops.def(
       "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
       "-> ()");
@@ -582,7 +624,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "bool skip_weighted, "
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
-#endif
+#endif  // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) &&
+        // !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__)
   ops.def(
       "mla_decode_kvcache("
       "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -595,6 +652,30 @@ 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

csrc/custom_all_gather_reduce_scatter.cuh

@@ -0,0 +1,327 @@
+#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

csrc/custom_all_reduce.cuh

@@ -1,299 +1,8 @@
 #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>
+#include "custom_collective_common.cuh"
 
 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)
@@ -616,6 +325,21 @@ 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));
@@ -625,8 +349,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

csrc/custom_collective_common.cuh

@@ -0,0 +1,332 @@
+#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

csrc/flashkda_registration.cpp

@@ -0,0 +1,17 @@
+#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)

csrc/fs_io.cpp

@@ -3,19 +3,155 @@
 
 #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" {
 
-static void _batch_lookup(const std::vector<const char*>& paths,
+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,
                           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.
@@ -51,11 +187,157 @@ 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},
 };
 

csrc/libtorch_stable/activation_kernels.cu

@@ -464,6 +464,66 @@ __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) \
@@ -553,6 +613,54 @@ 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.

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -21,7 +21,10 @@ __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_num_tokens,
+    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 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.
@@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel(
       }
     }
     if (output_lse != nullptr && pack_idx == 0) {
-      float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
-      output_lse[head_idx * num_tokens + token_idx] = s_lse;
+      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;
     }
     return;
   }
 
   // For tokens within prefix range, merge prefix and suffix
-  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
-  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
+  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];
   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;
 
@@ -132,7 +139,8 @@ __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 * num_tokens + token_idx] = max_lse;
+      output_lse[head_idx * output_lse_head_stride +
+                 token_idx * output_lse_token_stride] = max_lse;
     }
     return;
   }
@@ -187,7 +195,8 @@ __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 * num_tokens + token_idx] = out_lse;
+    output_lse[head_idx * output_lse_head_stride +
+               token_idx * output_lse_token_stride] = out_lse;
   }
 }
 
@@ -221,6 +230,9 @@ __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);                           \
   }
 
@@ -259,6 +271,19 @@ 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,

csrc/libtorch_stable/cache_kernels.cu

@@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel(
   copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
 }
 
+// Grouped variant of concat_and_cache_mla: inserts the context K/V for every
+// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
+// reads its own cache base pointer from kv_cache_ptrs (same pointer-array
+// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no
+// scaling or quantization; scalar_t is uint16_t for portability.
+template <typename scalar_t>
+__global__ void concat_and_cache_mla_grouped_kernel(
+    const scalar_t* __restrict__ kv_c,  // [num_layers, num_tokens,
+                                        // kv_lora_rank]
+    const scalar_t* __restrict__ k_pe,  // [num_layers, num_tokens, pe_dim]
+    const int64_t* __restrict__ kv_cache_ptrs,  // [num_layers]
+    const int64_t* __restrict__ slot_mapping,   // [num_layers, num_tokens]
+    const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
+    const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
+    const int64_t slot_layer_stride, const int64_t block_stride,
+    const int64_t entry_stride, const int kv_lora_rank, const int pe_dim,
+    const int block_size) {
+  const int64_t token_idx = blockIdx.x;
+  const int64_t layer_idx = blockIdx.y;
+  const int64_t slot_idx =
+      slot_mapping[layer_idx * slot_layer_stride + token_idx];
+  // NOTE: slot_idx can be -1 if the token is padded
+  if (slot_idx < 0) {
+    return;
+  }
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+
+  scalar_t* __restrict__ kv_cache =
+      reinterpret_cast<scalar_t*>(kv_cache_ptrs[layer_idx]);
+  const scalar_t* __restrict__ kv_c_layer =
+      kv_c + layer_idx * kv_c_layer_stride;
+  const scalar_t* __restrict__ k_pe_layer =
+      k_pe + layer_idx * k_pe_layer_stride;
+
+  auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
+                  int size, int offset) {
+    for (int i = threadIdx.x; i < size; i += blockDim.x) {
+      const int64_t src_idx = token_idx * src_token_stride + i;
+      const int64_t dst_idx =
+          block_idx * block_stride + block_offset * entry_stride + i + offset;
+      kv_cache[dst_idx] = src[src_idx];
+    }
+  };
+
+  copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0);
+  copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank);
+}
+
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void concat_and_cache_ds_mla_kernel(
     const scalar_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
@@ -902,6 +951,53 @@ void concat_and_cache_mla(
   }
 }
 
+void concat_and_cache_mla_grouped(
+    torch::stable::Tensor& kv_c,  // [num_layers, num_tokens, kv_lora_rank]
+    torch::stable::Tensor& k_pe,  // [num_layers, num_tokens, pe_dim]
+    torch::stable::Tensor& kv_cache_ptrs,  // [num_layers] int64, on device
+    torch::stable::Tensor& slot_mapping,   // [num_layers, num_tokens] int64
+    int64_t block_size, int64_t block_stride, int64_t entry_stride) {
+  int num_layers = kv_c.size(0);
+  int num_tokens = kv_c.size(1);
+  int kv_lora_rank = kv_c.size(2);
+  int pe_dim = k_pe.size(2);
+
+  STD_TORCH_CHECK(
+      kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
+          k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=",
+      kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
+  STD_TORCH_CHECK(
+      kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
+      "kv_cache_ptrs must be int64");
+
+  if (num_tokens == 0 || num_layers == 0) {
+    return;
+  }
+
+  const int64_t kv_c_layer_stride = kv_c.stride(0);
+  const int64_t kv_c_token_stride = kv_c.stride(1);
+  const int64_t k_pe_layer_stride = k_pe.stride(0);
+  const int64_t k_pe_token_stride = k_pe.stride(1);
+  const int64_t slot_layer_stride = slot_mapping.stride(0);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      kv_c.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  dim3 grid(num_tokens, num_layers);
+  dim3 block(std::min(kv_lora_rank, 512));
+  vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
+      <<<grid, block, 0, stream>>>(
+          reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
+          reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
+          kv_cache_ptrs.const_data_ptr<int64_t>(),
+          slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
+          kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
+          slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
+          block_size);
+}
+
 namespace vllm {
 
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
@@ -1025,6 +1121,9 @@ __global__ void gather_and_maybe_dequant_cache(
     batch_offset += offset;
     int32_t block_table_id = batch_offset / block_size;
     int32_t slot_id = batch_offset % block_size;
+    // seq_starts may push the block index past the end of the batch's block
+    // table row.
+    if (block_table_id >= block_table_stride) continue;
     int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
     int32_t block_id = block_table[block_table_offset];
     int64_t cache_offset =

csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu

@@ -0,0 +1,362 @@
+#include "torch_utils.h"
+
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "custom_all_reduce.cuh"
+#include "custom_all_gather_reduce_scatter.cuh"
+
+namespace vllm {
+
+void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output,
+                                int size_bytes, int threads, int block_limit) {
+  if (size_bytes % sizeof(CopyPack) != 0)
+    throw std::runtime_error(
+        "custom allgather requires input byte size to be a multiple of " +
+        std::to_string(sizeof(CopyPack)));
+
+  auto ptrs = buffers_.at(input);
+  int size_per_rank = size_bytes / sizeof(CopyPack);
+  int total_size = size_per_rank * world_size_;
+  int blocks = std::min(block_limit, (total_size + threads - 1) / threads);
+
+#define AG_CASE(ngpus)                                                   \
+  case ngpus:                                                            \
+    cross_device_all_gather<ngpus><<<blocks, threads, 0, stream>>>(      \
+        ptrs, sg_, self_sg_, reinterpret_cast<CopyPack*>(output), rank_, \
+        size_per_rank);                                                  \
+    break;
+
+  switch (world_size_) {
+    AG_CASE(2)
+    AG_CASE(4)
+    AG_CASE(6)
+    AG_CASE(8)
+    default:
+      throw std::runtime_error(
+          "custom allgather only supports num gpus in (2,4,6,8)");
+  }
+#undef AG_CASE
+}
+
+template <typename T>
+void CustomAllreduce::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) {
+  if (size_bytes % sizeof(typename packed_t<T>::P) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
+    throw std::runtime_error(
+        "MNNVL Lamport allgather requires 16-byte aligned sizes");
+
+  auto ptrs = buffers_.at(local_buffer);
+  int size_per_rank = size_bytes / sizeof(typename packed_t<T>::P);
+  int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
+  int blocks =
+      (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads;
+
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
+  cudaLaunchAttribute attributes[1]{};
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attributes[0].val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{.gridDim = dim3(blocks),
+                            .blockDim = dim3(kMnnvlLamportAgThreads),
+                            .dynamicSmemBytes = 0,
+                            .stream = stream,
+                            .attrs = attributes,
+                            .numAttrs = 1};
+  #define MNNVL_LAMPORT_AG_LAUNCH(ngpus)                                       \
+    CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather<T, ngpus>, \
+                                 ptrs, input, output,                          \
+                                 reinterpret_cast<T*>(multicast_buffer),       \
+                                 epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_AG_LAUNCH(ngpus)                                 \
+    mnnvl_lamport_all_gather<T, ngpus>                                   \
+        <<<blocks, kMnnvlLamportAgThreads, 0, stream>>>(                 \
+            ptrs, input, output, reinterpret_cast<T*>(multicast_buffer), \
+            epochs, rank_, size_per_rank, stage_size)
+#endif
+
+#define MNNVL_LAMPORT_AG_CASE(ngpus) \
+  case ngpus:                        \
+    MNNVL_LAMPORT_AG_LAUNCH(ngpus);  \
+    break;
+
+  switch (world_size_) {
+    MNNVL_LAMPORT_AG_CASE(2)
+    MNNVL_LAMPORT_AG_CASE(4)
+    MNNVL_LAMPORT_AG_CASE(6)
+    MNNVL_LAMPORT_AG_CASE(8)
+    MNNVL_LAMPORT_AG_CASE(16)
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)");
+  }
+#undef MNNVL_LAMPORT_AG_CASE
+#undef MNNVL_LAMPORT_AG_LAUNCH
+}
+
+template <typename T>
+void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output,
+                                     int size, int threads, int block_limit) {
+  auto packed_size = packed_t<T>::P::size;
+  if (size % (packed_size * world_size_) != 0)
+    throw std::runtime_error(
+        "custom reduce-scatter requires each output shard byte size to be "
+        "a multiple of 16");
+
+  auto ptrs = buffers_.at(input);
+  int size_per_rank = size / packed_size / world_size_;
+  int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads);
+
+#define RS_CASE(ngpus)                                                     \
+  case ngpus:                                                              \
+    cross_device_reduce_scatter<T, ngpus><<<blocks, threads, 0, stream>>>( \
+        ptrs, sg_, self_sg_, output, rank_, size_per_rank);                \
+    break;
+
+  switch (world_size_) {
+    RS_CASE(2)
+    RS_CASE(4)
+    RS_CASE(6)
+    RS_CASE(8)
+    default:
+      throw std::runtime_error(
+          "custom reduce-scatter only supports num gpus in (2,4,6,8)");
+  }
+#undef RS_CASE
+}
+
+template <typename T>
+void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream,
+                                                   T* input, T* output,
+                                                   void* local_buffer,
+                                                   uint32_t* epochs, int size,
+                                                   int stage_size_bytes) {
+  auto packed_size = packed_t<T>::P::size;
+  if (size % (packed_size * world_size_) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
+    throw std::runtime_error(
+        "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes");
+
+  auto ptrs = buffers_.at(local_buffer);
+  int size_per_rank = size / packed_size / world_size_;
+  int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
+  int blocks_per_rank =
+      (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads;
+  int blocks = blocks_per_rank * world_size_;
+
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
+  cudaLaunchAttribute attributes[1]{};
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attributes[0].val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{.gridDim = dim3(blocks),
+                            .blockDim = dim3(kMnnvlLamportRsThreads),
+                            .dynamicSmemBytes = 0,
+                            .stream = stream,
+                            .attrs = attributes,
+                            .numAttrs = 1};
+  #define MNNVL_LAMPORT_RS_LAUNCH(ngpus)                                      \
+    CUDACHECK(cudaLaunchKernelEx(                                             \
+        &config, &mnnvl_lamport_reduce_scatter_kernel<T, ngpus>, ptrs, input, \
+        output, epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_RS_LAUNCH(ngpus)                 \
+    mnnvl_lamport_reduce_scatter_kernel<T, ngpus>        \
+        <<<blocks, kMnnvlLamportRsThreads, 0, stream>>>( \
+            ptrs, input, output, epochs, rank_, size_per_rank, stage_size)
+#endif
+
+#define MNNVL_LAMPORT_RS_CASE(ngpus) \
+  case ngpus:                        \
+    MNNVL_LAMPORT_RS_LAUNCH(ngpus);  \
+    break;
+
+  switch (world_size_) {
+    MNNVL_LAMPORT_RS_CASE(2)
+    MNNVL_LAMPORT_RS_CASE(4)
+    MNNVL_LAMPORT_RS_CASE(6)
+    MNNVL_LAMPORT_RS_CASE(8)
+    MNNVL_LAMPORT_RS_CASE(16)
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport reduce-scatter only supports num gpus in "
+          "(2,4,6,8,16)");
+  }
+#undef MNNVL_LAMPORT_RS_CASE
+#undef MNNVL_LAMPORT_RS_LAUNCH
+}
+
+}  // namespace vllm
+
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+bool _is_weak_contiguous(torch::stable::Tensor& t);
+
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                       torch::stable::Tensor& out, fptr_t _reg_buffer,
+                       int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  STD_TORCH_CHECK(reg_buffer != nullptr);
+  STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
+  STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
+                                 cudaMemcpyDeviceToDevice, stream));
+  fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size);
+}
+
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                              torch::stable::Tensor& out, fptr_t _local_buffer,
+                              fptr_t _multicast_buffer, fptr_t _epoch_buffer,
+                              int64_t stage_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes);
+  auto local_buffer = reinterpret_cast<void*>(_local_buffer);
+  auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_allgather<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_allgather<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->mnnvl_lamport_allgather<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport allgather only supports float32, float16 and "
+          "bfloat16");
+  }
+}
+
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                           torch::stable::Tensor& out, fptr_t _reg_buffer,
+                           int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  STD_TORCH_CHECK(reg_buffer != nullptr);
+  STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
+  STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
+                                 cudaMemcpyDeviceToDevice, stream));
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(reg_buffer),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), inp.numel());
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                               reinterpret_cast<half*>(out.mutable_data_ptr()),
+                               inp.numel());
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), inp.numel());
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "custom reduce-scatter only supports float32, float16 and bfloat16");
+  }
+}
+
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                  torch::stable::Tensor& out,
+                                  fptr_t _local_buffer, fptr_t _epoch_buffer,
+                                  int64_t stage_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  STD_TORCH_CHECK(input_size <= stage_sz_bytes);
+  auto local_buffer = reinterpret_cast<void*>(_local_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          epochs, inp.numel(), stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_reduce_scatter<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer, epochs,
+          inp.numel(), stage_sz_bytes);
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->mnnvl_lamport_reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
+          epochs, inp.numel(), stage_sz_bytes);
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport reduce-scatter only supports float32, float16 and "
+          "bfloat16");
+  }
+}

csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp

@@ -0,0 +1,29 @@
+#include "ops.h"
+#include "core/registration.h"
+
+#include <torch/csrc/stable/library.h>
+
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) {
+  custom_ag_rs.def(
+      "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int "
+      "local_buffer, int multicast_buffer, int epoch_buffer, int "
+      "stage_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
+      "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
+  custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather));
+  custom_ag_rs.impl("mnnvl_lamport_all_gather",
+                    TORCH_BOX(&mnnvl_lamport_all_gather));
+  custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter));
+  custom_ag_rs.impl("mnnvl_lamport_reduce_scatter",
+                    TORCH_BOX(&mnnvl_lamport_reduce_scatter));
+}

csrc/libtorch_stable/custom_all_reduce.cu

@@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
                       torch::stable::Tensor& rank_data, int64_t rank,
                       bool fully_connected) {
   int world_size = fake_ipc_ptrs.size();
-  if (world_size > 8)
-    throw std::invalid_argument("world size > 8 is not supported");
+  if (world_size > vllm::kMaxCustomCollectiveRanks)
+    throw std::invalid_argument("world size > 16 is not supported");
   if (world_size % 2 != 0)
     throw std::invalid_argument("Odd num gpus is not supported for now");
   if (rank < 0 || rank >= world_size)
     throw std::invalid_argument("invalid rank passed in");
 
-  vllm::Signal* ipc_ptrs[8];
+  vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < world_size; i++) {
     ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
   }
@@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
   STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
-  void* ipc_ptrs[8];
+  void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
     ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
   }

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

@@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
 #endif
 }
 
-template <typename T, int kHdIn, int kHdOut, int kTileN>
+template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
-                      cudaStream_t const stream) {
-  constexpr int gemm_m = kHdOut;  // 2112
-  int const gemm_n = num_tokens;  // 1-16
-  constexpr int gemm_k = kHdIn;   // 7168
+                      cudaStream_t const stream, bool enable_pdl) {
+  constexpr int gemm_m = kHdOut;
+  int const gemm_n = num_tokens;
+  constexpr int gemm_k = kHdIn;
   constexpr int batch_size = 1;
   std::swap(mat_a, mat_b);
   constexpr int tile_m = 16;
-  constexpr int tile_n = kTileN;                        // 8 or 16
-  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+  constexpr int tile_n = kTileN;
+  constexpr int tile_k = kTileK;
   constexpr int max_stage_cnt =
       1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
   constexpr int k_iter_cnt = gemm_k / tile_k;
@@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  attrs[0].val.programmaticStreamSerializationAllowed =
+      enable_pdl || getEnvEnablePDL();
   config.numAttrs = 1;
   config.attrs = attrs;
   if (smem_bytes >= (48 * 1024)) {
@@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
                      output, mat_a, mat_b, gemm_n);
 }
 
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
-
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
+template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
+void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
+                               int num_tokens, cudaStream_t const stream,
+                               bool enable_pdl) {
+  if (num_tokens <= 8) {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  } else {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  }
+}
 
 void dsv3_fused_a_gemm(torch::stable::Tensor& output,
                        torch::stable::Tensor const& mat_a,
-                       torch::stable::Tensor const& mat_b) {
+                       torch::stable::Tensor const& mat_b, bool enable_pdl) {
   STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
   int const num_tokens = mat_a.size(0);
   int const hd_in = mat_a.size(1);
   int const hd_out = mat_b.size(1);
 
-  constexpr int kHdIn = 7168;
-  constexpr int kHdOut = 2112;
   STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                   "required 1 <= mat_a.shape[0] <= 16");
-  STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
-  STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
   STD_TORCH_CHECK(output.size(0) == num_tokens,
                   "required output.shape[0] == mat_a.shape[0]");
   STD_TORCH_CHECK(output.size(1) == hd_out,
                   "required output.shape[1] == mat_b.shape[1]");
+  STD_TORCH_CHECK(mat_b.size(0) == hd_in,
+                  "required mat_b.shape[0] == mat_a.shape[1]");
 
-  STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
-  STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
-  STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+  STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() &&
+                      mat_a.get_device_index() == output.get_device_index(),
+                  "mat_a, mat_b, and output must be on the same device");
+
+  // The kernels index global memory with raw pointers and packed strides, so
+  // reject any padded or transposed view rather than reading out of bounds.
+  STD_TORCH_CHECK(
+      mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
+      "mat_a must be a packed row-major [num_tokens, hd_in] tensor");
+  STD_TORCH_CHECK(
+      output.stride(0) == hd_out && output.stride(1) == 1,
+      "output must be a packed row-major [num_tokens, hd_out] tensor");
+  STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in,
+                  "mat_b must be a packed column-major [hd_in, hd_out] tensor");
 
   STD_TORCH_CHECK(
       mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
@@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
   STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
 
   auto stream = get_current_cuda_stream(mat_a.get_device_index());
-  if (num_tokens <= 8) {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
-  } else {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
+  auto* output_ptr =
+      reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+  auto const* mat_a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  auto const* mat_b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
+#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT)                                 \
+  if (hd_in == HD_IN && hd_out == HD_OUT) {                                \
+    invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>(               \
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \
+    return;                                                                \
   }
+
+  // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
+  // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection.
+  DISPATCH_DSV3_SHAPE(7168, 1536)
+  DISPATCH_DSV3_SHAPE(7168, 2112)
+  DISPATCH_DSV3_SHAPE(1536, 2304)
+  DISPATCH_DSV3_SHAPE(1536, 4608)
+  DISPATCH_DSV3_SHAPE(7168, 3584)
+  DISPATCH_DSV3_SHAPE(768, 7168)
+  // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent
+  // because hd_in=2112 is not a multiple of any supported tile_k.
+  DISPATCH_DSV3_SHAPE(1536, 1152)
+  DISPATCH_DSV3_SHAPE(7168, 768)
+  DISPATCH_DSV3_SHAPE(7168, 3216)
+  DISPATCH_DSV3_SHAPE(7168, 4224)
+
+#ifdef VLLM_K3_BENCH_SHAPES
+  // The selector routes these shapes to CuTe or the default GEMM, so they are
+  // never reached in production. They are compiled only for offline
+  // DSV3-vs-CuTe benchmarking.
+  DISPATCH_DSV3_SHAPE(7168, 6288)
+  DISPATCH_DSV3_SHAPE(1536, 7168)
+  DISPATCH_DSV3_SHAPE(3584, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 8448)
+  DISPATCH_DSV3_SHAPE(7168, 20480)
+  DISPATCH_DSV3_SHAPE(7168, 3072)
+  DISPATCH_DSV3_SHAPE(7168, 12448)
+  DISPATCH_DSV3_SHAPE(3072, 7168)
+  DISPATCH_DSV3_SHAPE(8448, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 16896)
+  DISPATCH_DSV3_SHAPE(7168, 40960)
+#endif
+
+#undef DISPATCH_DSV3_SHAPE
+
+  if (hd_in == 128 && hd_out == 1536) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 128 && hd_out == 3072) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple
+  // of 256, so both need the 128 tile_k.
+  if (hd_in == 128 && hd_out == 768) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 384 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#ifdef VLLM_K3_BENCH_SHAPES
+  if (hd_in == 4224 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#endif
+
+  STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {