Merge branch 'main' into wentao-fix-mypy-models-a-b

Signed-off-by: johnny <johnnyychiu@gmail.com>

.buildkite/check-torch-abi.py

@@ -0,0 +1,102 @@
+# 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, ...] = (
+    "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,8 @@ 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 compressed tensors FP8 test"
     depends_on:
       - image-build-xpu
@@ -168,4 +169,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/release-pipeline.yaml

@@ -31,8 +31,46 @@ steps:
       - text: "What is the release version?"
         key: release-version
 
-  - group: "Build Python wheels"
+  - group: "Build CUDA 13.0 Python wheels"
     key: "build-wheels"
+    steps:
+      - label: "Build wheel - aarch64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-arm64-cuda-13-0
+        agents:
+          queue: arm64_cpu_queue_release
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+        env:
+          DOCKER_BUILDKIT: "1"
+
+      - label: "Build wheel - x86_64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-x86-cuda-13-0
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+        env:
+          DOCKER_BUILDKIT: "1"
+
+  - block: "Unblock to build additional Python wheels"
+    depends_on: ~
+    key: block-build-additional-wheels
+    if: build.env("NIGHTLY") != "1"
+
+  - group: "Build additional Python wheels"
+    key: "build-additional-wheels"
+    depends_on: block-build-additional-wheels
+    allow_dependency_failure: true
     steps:
       - label: "Build wheel - aarch64 - CUDA 12.9"
         depends_on: ~
@@ -48,20 +86,6 @@ steps:
         env:
           DOCKER_BUILDKIT: "1"
 
-      - label: "Build wheel - aarch64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-arm64-cuda-13-0
-        agents:
-          queue: arm64_cpu_queue_release
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
-        env:
-          DOCKER_BUILDKIT: "1"
-
       - label: "Build wheel - aarch64 - CPU"
         depends_on: ~
         id: build-wheel-arm64-cpu
@@ -113,7 +137,7 @@ steps:
           - 'mv artifacts/reassembled/wheel "artifacts/dist/$$wheel_name"'
           - "aws sts get-caller-identity"
           - "VLLM_WHEEL_PLATFORM=macos bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
         plugins:
           - aws-assume-role-with-web-identity#v1.6.0:
               role-arn: arn:aws:iam::936637512419:role/vllm-release-macos-wheel-uploader
@@ -133,20 +157,6 @@ steps:
         env:
           DOCKER_BUILDKIT: "1"
 
-      - label: "Build wheel - x86_64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-x86-cuda-13-0
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
-        env:
-          DOCKER_BUILDKIT: "1"
-
       - label: "Build wheel - x86_64 - CPU"
         depends_on: ~
         id: build-wheel-x86-cpu
@@ -162,12 +172,26 @@ steps:
           DOCKER_BUILDKIT: "1"
 
   - label: "Generate and upload wheel indices"
+    key: generate-wheel-indices
     depends_on: "build-wheels"
     allow_dependency_failure: true
+    if: build.env("NIGHTLY") != "1"
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "UPDATE_VERSION_INDEX=0 bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+
+  - label: "Regenerate indices with additional wheels"
+    key: generate-additional-wheel-indices
+    depends_on:
+      - build-wheels
+      - build-additional-wheels
+      - generate-wheel-indices
+    allow_dependency_failure: true
     agents:
       queue: cpu_queue_release
     commands:
-      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+      - 'UPDATE_NIGHTLY_INDEX="$${NIGHTLY:-0}" bash .buildkite/scripts/generate-and-upload-nightly-index.sh'
 
   - block: "Unblock to build release Docker images"
     depends_on: ~
@@ -566,366 +590,370 @@ steps:
   #
   # =============================================================================
 
-  # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
-  - label: ":rocm: Build ROCm Base Image & Wheels"
-    id: build-rocm-base-wheels
+  - group: "Build ROCm Wheel / Image "
+    key: "build-rocm-wheel-image"
     depends_on: ~
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - |
-        set -euo pipefail
+    steps:
+      # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
+      - label: ":rocm: Build ROCm Base Image & Wheels"
+        id: build-rocm-base-wheels
+        depends_on: ~
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - |
+            set -euo pipefail
 
-        # Generate cache key
-        CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-        ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
+            # Generate cache key
+            CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+            ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
 
-        echo "========================================"
-        echo "ROCm Base Build Configuration"
-        echo "========================================"
-        echo "  CACHE_KEY: $${CACHE_KEY}"
-        echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
-        echo "========================================"
-        
-        # Login to ECR
-        aws ecr-public get-login-password --region us-east-1 | \
-          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
-        IMAGE_EXISTS=false
-        WHEELS_EXIST=false
-        
-        # Check ECR for Docker image
+            echo "========================================"
+            echo "ROCm Base Build Configuration"
+            echo "========================================"
+            echo "  CACHE_KEY: $${CACHE_KEY}"
+            echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
+            echo "========================================"
+            
+            # Login to ECR
+            aws ecr-public get-login-password --region us-east-1 | \
+              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+            
+            IMAGE_EXISTS=false
+            WHEELS_EXIST=false
+            
+            # Check ECR for Docker image
 
-        if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
-          IMAGE_EXISTS=true
-          echo "ECR image cache HIT"
-        fi
-        
-        # Check S3 for wheels
-        WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
-        if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
-          WHEELS_EXIST=true
-          echo "S3 wheels cache HIT"
-        fi
+            if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
+              IMAGE_EXISTS=true
+              echo "ECR image cache HIT"
+            fi
+            
+            # Check S3 for wheels
+            WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
+            if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
+              WHEELS_EXIST=true
+              echo "S3 wheels cache HIT"
+            fi
 
-        
-        # Scenario 1: Both cached (best case)
-        if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
-          echo ""
-          echo "FULL CACHE HIT - Reusing both image and wheels"
-          echo ""
+            
+            # Scenario 1: Both cached (best case)
+            if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
+              echo ""
+              echo "FULL CACHE HIT - Reusing both image and wheels"
+              echo ""
 
-          # Download wheels
-          .buildkite/scripts/cache-rocm-base-wheels.sh download
-          
-          # Save ECR tag for downstream jobs
-          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-          
-        # Scenario 2: Full rebuild needed
-        else
-          echo ""
-          echo " CACHE MISS - Building from scratch..."
-          echo ""
-          
-          # Build full base image and push to ECR
-          DOCKER_BUILDKIT=1 docker buildx build \
-            --file docker/Dockerfile.rocm_base \
-            --tag "$${ECR_CACHE_TAG}" \
-            --build-arg USE_SCCACHE=1 \
-            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-            --push \
-            .
-          
-          # Build wheel extraction stage
-          DOCKER_BUILDKIT=1 docker buildx build \
-            --file docker/Dockerfile.rocm_base \
-            --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
-            --target debs_wheel_release \
-            --build-arg USE_SCCACHE=1 \
-            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-            --load \
-            .
-          
-          # Extract and upload wheels
-          mkdir -p artifacts/rocm-base-wheels
-          cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
-          docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
-          docker rm $${cid}
-          
-          .buildkite/scripts/cache-rocm-base-wheels.sh upload
+              # Download wheels
+              .buildkite/scripts/cache-rocm-base-wheels.sh download
+              
+              # Save ECR tag for downstream jobs
+              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+              
+            # Scenario 2: Full rebuild needed
+            else
+              echo ""
+              echo " CACHE MISS - Building from scratch..."
+              echo ""
+              
+              # Build full base image and push to ECR
+              DOCKER_BUILDKIT=1 docker buildx build \
+                --file docker/Dockerfile.rocm_base \
+                --tag "$${ECR_CACHE_TAG}" \
+                --build-arg USE_SCCACHE=1 \
+                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+                --build-arg SCCACHE_REGION_NAME=us-west-2 \
+                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+                --push \
+                .
+              
+              # Build wheel extraction stage
+              DOCKER_BUILDKIT=1 docker buildx build \
+                --file docker/Dockerfile.rocm_base \
+                --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
+                --target debs_wheel_release \
+                --build-arg USE_SCCACHE=1 \
+                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+                --build-arg SCCACHE_REGION_NAME=us-west-2 \
+                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+                --load \
+                .
+              
+              # Extract and upload wheels
+              mkdir -p artifacts/rocm-base-wheels
+              cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
+              docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
+              docker rm $${cid}
+              
+              .buildkite/scripts/cache-rocm-base-wheels.sh upload
 
-          # Cache base docker image to ECR
-          docker push "$${ECR_CACHE_TAG}"
-          
-          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-          
-          echo ""
-          echo " Build complete - Image and wheels cached"
-        fi
+              # Cache base docker image to ECR
+              docker push "$${ECR_CACHE_TAG}"
+              
+              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+              
+              echo ""
+              echo " Build complete - Image and wheels cached"
+            fi
 
-    artifact_paths:
-      - "artifacts/rocm-base-wheels/*.whl"
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
+        artifact_paths:
+          - "artifacts/rocm-base-wheels/*.whl"
+        env:
+          DOCKER_BUILDKIT: "1"
+          S3_BUCKET: "vllm-wheels"
 
-  # ROCm Job 2: Build vLLM ROCm Wheel
-  - label: ":python: Build vLLM ROCm Wheel - x86_64"
-    id: build-rocm-vllm-wheel
-    depends_on:
-      - step: build-rocm-base-wheels
-        allow_failure: false
-    agents:
-      queue: cpu_queue_release
-    timeout_in_minutes: 180
-    commands:
-      # Download artifacts and prepare Docker image
-      - |
-        set -euo pipefail
+      # ROCm Job 2: Build vLLM ROCm Wheel
+      - label: ":python: Build vLLM ROCm Wheel - x86_64"
+        id: build-rocm-vllm-wheel
+        depends_on:
+          - step: build-rocm-base-wheels
+            allow_failure: false
+        agents:
+          queue: cpu_queue_release
+        timeout_in_minutes: 180
+        commands:
+          # Download artifacts and prepare Docker image
+          - |
+            set -euo pipefail
 
-        # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
-        # This fixes version detection when tags are moved/force-pushed
-        echo "Fetching latest tags from origin..."
-        git fetch --tags --force origin
-        
-        # Log tag information for debugging version detection
-        echo "========================================"
-        echo "Git Tag Verification"
-        echo "========================================"
-        echo "Current HEAD: $(git rev-parse HEAD)"
-        echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
-        echo ""
-        echo "Recent tags (pointing to commits near HEAD):"
-        git tag -l --sort=-creatordate | head -5
-        echo "setuptools_scm version detection:"
-        pip install -q setuptools_scm 2>/dev/null || true
-        python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
-        echo "========================================"
+            # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
+            # This fixes version detection when tags are moved/force-pushed
+            echo "Fetching latest tags from origin..."
+            git fetch --tags --force origin
+            
+            # Log tag information for debugging version detection
+            echo "========================================"
+            echo "Git Tag Verification"
+            echo "========================================"
+            echo "Current HEAD: $(git rev-parse HEAD)"
+            echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
+            echo ""
+            echo "Recent tags (pointing to commits near HEAD):"
+            git tag -l --sort=-creatordate | head -5
+            echo "setuptools_scm version detection:"
+            pip install -q setuptools_scm 2>/dev/null || true
+            python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
+            echo "========================================"
 
-        # Download wheel artifacts from current build
-        echo "Downloading wheel artifacts from current build"
-        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+            # Download wheel artifacts from current build
+            echo "Downloading wheel artifacts from current build"
+            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
 
-        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-        if [ -z "$${ECR_IMAGE_TAG}" ]; then
-          echo "ERROR: rocm-base-image-tag metadata not found"
-          echo "This should have been set by the build-rocm-base-wheels job"
-          exit 1
-        fi
-        
-        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-        
-        # Login to ECR
-        aws ecr-public get-login-password --region us-east-1 | \
-          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
-        # Pull base Docker image from ECR
-        docker pull "$${ECR_IMAGE_TAG}"
-        
-        echo "Loaded base image: $${ECR_IMAGE_TAG}"
-        
-        # Prepare base wheels for Docker build context
-        mkdir -p docker/context/base-wheels
-        touch docker/context/base-wheels/.keep
-        cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
-        echo "Base wheels for vLLM build:"
-        ls -lh docker/context/base-wheels/
+            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+            if [ -z "$${ECR_IMAGE_TAG}" ]; then
+              echo "ERROR: rocm-base-image-tag metadata not found"
+              echo "This should have been set by the build-rocm-base-wheels job"
+              exit 1
+            fi
+            
+            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+            
+            # Login to ECR
+            aws ecr-public get-login-password --region us-east-1 | \
+              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+            
+            # Pull base Docker image from ECR
+            docker pull "$${ECR_IMAGE_TAG}"
+            
+            echo "Loaded base image: $${ECR_IMAGE_TAG}"
+            
+            # Prepare base wheels for Docker build context
+            mkdir -p docker/context/base-wheels
+            touch docker/context/base-wheels/.keep
+            cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
+            echo "Base wheels for vLLM build:"
+            ls -lh docker/context/base-wheels/
 
-        echo "========================================"
-        echo "Building vLLM wheel with:"
-        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-        echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
-        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-        echo "========================================"
+            echo "========================================"
+            echo "Building vLLM wheel with:"
+            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+            echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
+            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+            echo "========================================"
 
-        # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
-        DOCKER_BUILDKIT=1 docker build \
-          --file docker/Dockerfile.rocm \
-          --target export_vllm_wheel_release \
-          --output type=local,dest=rocm-dist \
-          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-          --build-arg REMOTE_VLLM=0 \
-          --build-arg GIT_REPO_CHECK=1 \
-          --build-arg USE_SCCACHE=1 \
-          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-          --build-arg SCCACHE_REGION_NAME=us-west-2 \
-          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-          .
-        echo "Built vLLM wheel:"
-        ls -lh rocm-dist/*.whl
-        # Copy wheel to artifacts directory
-        mkdir -p artifacts/rocm-vllm-wheel
-        cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
-        echo "Final vLLM wheel:"
-        ls -lh artifacts/rocm-vllm-wheel/
-    artifact_paths:
-      - "artifacts/rocm-vllm-wheel/*.whl"
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
+            # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
+            DOCKER_BUILDKIT=1 docker build \
+              --file docker/Dockerfile.rocm \
+              --target export_vllm_wheel_release \
+              --output type=local,dest=rocm-dist \
+              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+              --build-arg REMOTE_VLLM=0 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+              --build-arg SCCACHE_REGION_NAME=us-west-2 \
+              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+              .
+            echo "Built vLLM wheel:"
+            ls -lh rocm-dist/*.whl
+            # Copy wheel to artifacts directory
+            mkdir -p artifacts/rocm-vllm-wheel
+            cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
+            echo "Final vLLM wheel:"
+            ls -lh artifacts/rocm-vllm-wheel/
+        artifact_paths:
+          - "artifacts/rocm-vllm-wheel/*.whl"
+        env:
+          DOCKER_BUILDKIT: "1"
+          S3_BUCKET: "vllm-wheels"
 
-  # ROCm Job 3: Upload Wheels to S3
-  - label: ":s3: Upload ROCm Wheels to S3"
-    id: upload-rocm-wheels
-    depends_on:
-      - step: build-rocm-vllm-wheel
-        allow_failure: false
-    agents:
-      queue: cpu_queue_release
-    timeout_in_minutes: 60
-    commands:
-      # Download all wheel artifacts and run upload
-      - |
-        set -euo pipefail
+      # ROCm Job 3: Upload Wheels to S3
+      - label: ":s3: Upload ROCm Wheels to S3"
+        id: upload-rocm-wheels
+        depends_on:
+          - step: build-rocm-vllm-wheel
+            allow_failure: false
+        agents:
+          queue: cpu_queue_release
+        timeout_in_minutes: 60
+        commands:
+          # Download all wheel artifacts and run upload
+          - |
+            set -euo pipefail
 
-        # Download artifacts from current build
-        echo "Downloading artifacts from current build"
-        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-        buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+            # Download artifacts from current build
+            echo "Downloading artifacts from current build"
+            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+            buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
 
-        # Run upload script
-        bash .buildkite/scripts/upload-rocm-wheels.sh
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
+            # # Run upload script
+            bash .buildkite/scripts/upload-rocm-wheels.sh
+        env:
+          DOCKER_BUILDKIT: "1"
+          S3_BUCKET: "vllm-wheels"
 
-  # ROCm Job 4: Annotate ROCm Wheel Release
-  - label: ":memo: Annotate ROCm wheel release"
-    id: annotate-rocm-release
-    depends_on:
-      - upload-rocm-wheels
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - "bash .buildkite/scripts/annotate-rocm-release.sh"
-    env:
-      S3_BUCKET: "vllm-wheels"
+      # ROCm Job 4: Annotate ROCm Wheel Release
+      - label: ":memo: Annotate ROCm wheel release"
+        id: annotate-rocm-release
+        depends_on:
+          - upload-rocm-wheels
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/annotate-rocm-release.sh"
+        env:
+          S3_BUCKET: "vllm-wheels"
 
-  # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
-  # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
-  # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
-  - block: "Generate Root Index for ROCm Wheels for Release"
-    key: block-generate-root-index-rocm-wheels
-    depends_on: upload-rocm-wheels
+      # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
+      # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
+      # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
+      - block: "Generate Root Index for ROCm Wheels for Release"
+        key: block-generate-root-index-rocm-wheels
+        depends_on: upload-rocm-wheels
 
-  - label: ":package: Generate Root Index for ROCm Wheels for Release"
-    depends_on: block-generate-root-index-rocm-wheels
-    id: generate-root-index-rocm-wheels
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
-    env:
-      S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm723"
+      - label: ":package: Generate Root Index for ROCm Wheels for Release"
+        depends_on: block-generate-root-index-rocm-wheels
+        id: generate-root-index-rocm-wheels
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
+        env:
+          S3_BUCKET: "vllm-wheels"
+          VARIANT: "rocm723"
 
-  # ROCm Job 6: Build ROCm Release Docker Image
-  - label: ":docker: Build release image - x86_64 - ROCm"
-    id: build-rocm-release-image
-    depends_on:
-      - step: block-build-release-images
-        allow_failure: true
-      - step: build-rocm-base-wheels
-        allow_failure: false
-    agents:
-      queue: cpu_queue_release
-    timeout_in_minutes: 60
-    commands:
-      - |
-        set -euo pipefail
-        
-        # Login to ECR
-        aws ecr-public get-login-password --region us-east-1 | \
-          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
-        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-        if [ -z "$${ECR_IMAGE_TAG}" ]; then
-          echo "ERROR: rocm-base-image-tag metadata not found"
-          echo "This should have been set by the build-rocm-base-wheels job"
-          exit 1
-        fi
-        
-        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-        
-        # Pull base Docker image from ECR
-        docker pull "$${ECR_IMAGE_TAG}"
-        
-        echo "Loaded base image: $${ECR_IMAGE_TAG}"
-        
-        # Pass the base image ECR tag to downstream steps (nightly publish)
-        buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
-        
-        echo "========================================"
-        echo "Building vLLM ROCm release image with:"
-        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-        echo "========================================"
-        
-        # Build vLLM ROCm release image using cached base
-        DOCKER_BUILDKIT=1 docker build \
-          --build-arg max_jobs=16 \
-          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-          --build-arg USE_SCCACHE=1 \
-          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-          --build-arg SCCACHE_REGION_NAME=us-west-2 \
-          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-          --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
-          --target vllm-openai \
-          --progress plain \
-          -f docker/Dockerfile.rocm .
-        
-        # Push to ECR
-        docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
+      # ROCm Job 6: Build ROCm Release Docker Image
+      - label: ":docker: Build release image - x86_64 - ROCm"
+        id: build-rocm-release-image
+        depends_on:
+          - step: block-build-release-images
+            allow_failure: true
+          - step: build-rocm-base-wheels
+            allow_failure: false
+        agents:
+          queue: cpu_queue_release
+        timeout_in_minutes: 60
+        commands:
+          - |
+            set -euo pipefail
+            
+            # Login to ECR
+            aws ecr-public get-login-password --region us-east-1 | \
+              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+            
+            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+            if [ -z "$${ECR_IMAGE_TAG}" ]; then
+              echo "ERROR: rocm-base-image-tag metadata not found"
+              echo "This should have been set by the build-rocm-base-wheels job"
+              exit 1
+            fi
+            
+            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+            
+            # Pull base Docker image from ECR
+            docker pull "$${ECR_IMAGE_TAG}"
+            
+            echo "Loaded base image: $${ECR_IMAGE_TAG}"
+            
+            # Pass the base image ECR tag to downstream steps (nightly publish)
+            buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
+            
+            echo "========================================"
+            echo "Building vLLM ROCm release image with:"
+            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+            echo "========================================"
+            
+            # Build vLLM ROCm release image using cached base
+            DOCKER_BUILDKIT=1 docker build \
+              --build-arg max_jobs=16 \
+              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+              --build-arg SCCACHE_REGION_NAME=us-west-2 \
+              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+              --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile.rocm .
+            
+            # Push to ECR
+            docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
 
-        echo ""
-        echo " Successfully built and pushed ROCm release image"
-        echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
-        echo ""
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
+            echo ""
+            echo " Successfully built and pushed ROCm release image"
+            echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
+            echo ""
+        env:
+          DOCKER_BUILDKIT: "1"
+          S3_BUCKET: "vllm-wheels"
 
-  - label: "Publish nightly XPU image to DockerHub"
-    depends_on:
-      - create-manifest-xpu
-    if: build.env("NIGHTLY") == "1"
-    agents:
-      queue: small_cpu_queue_release
-    commands:
-      - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
-      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - label: "Publish nightly XPU image to DockerHub"
+        depends_on:
+          - create-manifest-xpu
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
 
-  - label: "Publish nightly ROCm image to DockerHub"
-    depends_on:
-      - build-rocm-release-image
-    if: build.env("NIGHTLY") == "1"
-    agents:
-      queue: small_cpu_queue_release
-    commands:
-      - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
-      # Clean up old nightly builds (keep only last 14)
-      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
-      - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - label: "Publish nightly ROCm image to DockerHub"
+        depends_on:
+          - build-rocm-release-image
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
+          # Clean up old nightly builds (keep only last 14)
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
 
   # =============================================================================
   # Publish to DockerHub and PyPI (at the end so all builds complete first)
@@ -974,6 +1002,8 @@ steps:
         depends_on:
           - input-release-version
           - build-wheels
+          - build-additional-wheels
+          - generate-additional-wheel-indices
 
       - label: "Upload release wheels to PyPI"
         depends_on:

.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/generate-and-upload-nightly-index.sh

@@ -45,8 +45,10 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
 
-# copy to /nightly/ only if it is on the main branch and not a PR
-if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
+# copy to /nightly/ only when enabled for a main branch build that is not a PR
+if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
+      "$BUILDKITE_BRANCH" == "main" && \
+      "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
     echo "Uploading indices to overwrite /nightly/"
     aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
 fi
@@ -67,7 +69,7 @@ pure_version="${version%%+*}"
 echo "Pure version (without variant): $pure_version"
 
 # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "$version" != *"dev"* ]]; then
+if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
     echo "Re-generating indices for /$pure_version/"
     rm -rf "${INDICES_OUTPUT_DIR:?}"
     mkdir -p "$INDICES_OUTPUT_DIR"

.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/upload-rocm-wheels.sh

@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
 
-# Update rocm/nightly/ if on main branch and not a PR
-if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
+# Only scheduled nightly builds should update the moving nightly index.
+if [[ "${NIGHTLY:-0}" == "1" ]]; then
     echo "Updating rocm/nightly/ index..."
     aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
 fi
@@ -147,7 +147,7 @@ echo ""
 echo "Install command (by commit):"
 echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
 echo ""
-if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
+if [[ "${NIGHTLY:-0}" == "1" ]]; then
     echo "Install command (nightly):"
     echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
 fi

.buildkite/test-amd.yaml

@@ -40,7 +40,7 @@
 #####################################################################################################################################
 #                                                                                                                                   #
 # IMPORTANT:                                                                                                                        #
-#   * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are        #
+#   * Currently AMD CI has MI250 agents, MI300 agents, and MI355 agents. All upcoming feature improvements are                      #
 #         tracked in: https://github.com/vllm-project/vllm/issues/34994                                                             #
 #                                                                                                                                   #
 #-----------------------------------------------------------------------------------------------------------------------------------#
@@ -81,10 +81,8 @@
 #                                                     the above test.) Also run if model initialization test file is modified.      #
 #   * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model         #
 #                                                  source is modified, or when specified test files are modified.                   #
-#   * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to       #
-#                                          run plamo2 model in vLLM.                                                                #
-#   * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d)     #
-#                                                   and to run plamo2 model in vLLM.                                                #
+#   * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d).             #
+#   * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d).    #
 #   * [Multi-Modal Models (Standard) 1-4]:                                                                                          #
 #     - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function.            #
 #   * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock.                          #
@@ -171,20 +169,6 @@ steps:
   - pip install helion==1.1.0
   - pytest -v -s kernels/helion/
 
-- label: Kernels Mamba Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - csrc/mamba/
-  - tests/kernels/mamba
-  - vllm/model_executor/layers/mamba/ops
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s kernels/mamba
-
 #------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
@@ -207,6 +191,7 @@ steps:
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -240,6 +225,20 @@ steps:
   - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
   - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
 
+- label: Multi-Modal Processor (CPU) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  parallelism: 6
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/multimodal
+  - tests/models/registry.py
+  commands:
+  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #------------------------------------------------------------  mi250 · v1  -------------------------------------------------------------#
 
 - label: Batch Invariance (H100-MI250) # TBD
@@ -351,21 +350,21 @@ steps:
   commands:
   - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
 
-- label: V1 attention (H100-MI250) # TBD
+- label: V1 others (CPU) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/config/attention.py
-  - vllm/model_executor/layers/attention
-  - vllm/v1/attention
-  - tests/v1/attention
-  - vllm/_aiter_ops.py
-  - vllm/envs.py
-  - vllm/platforms/rocm.py
+  - vllm/
+  - tests/v1
   commands:
-  - pytest -v -s v1/attention
+  - 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 -m 'cpu_test' v1/kv_connector/unit
+  - pytest -v -s -m 'cpu_test' v1/metrics
 
 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#
 
@@ -408,6 +407,19 @@ steps:
   - pytest -v -s transformers_utils
   - pytest -v -s config
 
+- label: Python-only Installation # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/standalone_tests/python_only_compile.sh
+  - setup.py
+  - vllm/platforms/rocm.py
+  commands:
+  - bash standalone_tests/python_only_compile.sh
+
 #------------------------------------------------------------  mi250 · rust  -----------------------------------------------------------#
 
 - label: Rust Frontend Cargo Style + Clippy # TBD
@@ -445,6 +457,7 @@ steps:
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
   no_gpu: true
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - .buildkite/scripts/docker-build-metadata-args.sh
@@ -507,7 +520,7 @@ steps:
   - tests/models/
   commands:
   - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
+  - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
   - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
@@ -657,6 +670,30 @@ steps:
   - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
   - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
 
+- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/compilation/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_wrapper.py
+  - tests/entrypoints/llm/test_collective_rpc.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py
+
 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#
 
 - label: Platform Tests # TBD
@@ -679,6 +716,7 @@ steps:
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -870,6 +908,71 @@ steps:
   commands:
   - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
 
+- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/distributed/
+  - tests/v1/shutdown
+  - tests/v1/worker/test_worker_memory_snapshot.py
+  - vllm/platforms/rocm.py
+  commands:
+  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
+  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
+
+- label: Distributed Compile + Comm (4 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_4
+  num_gpus: 4
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/
+  - tests/distributed/test_pynccl
+  - tests/distributed/test_events
+  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/distributed/test_symm_mem_allreduce.py
+  - tests/distributed/test_multiproc_executor.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s distributed/test_pynccl.py
+  - pytest -v -s distributed/test_events.py
+  - pytest -v -s distributed/test_symm_mem_allreduce.py
+  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
+
+#----------------------------------------------------------  mi300 · engine  -----------------------------------------------------------#
+
+- label: Engine # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/engine
+  - tests/test_sequence
+  - tests/test_config
+  - tests/test_logger
+  - tests/test_vllm_port
+  commands:
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
+
 #--------------------------------------------------------  mi300 · entrypoints  --------------------------------------------------------#
 
 - label: Entrypoints Unit Tests # TBD
@@ -878,6 +981,7 @@ steps:
   dind: false
   agent_pool: mi300_1
   fast_check: true
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/entrypoints
@@ -982,6 +1086,7 @@ steps:
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  optional: true
   fast_check: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1010,6 +1115,7 @@ steps:
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  optional: true
   fast_check: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1024,6 +1130,7 @@ steps:
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  optional: true
   fast_check: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1345,6 +1452,27 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
 
+- label: LM Eval Large Models (4xH100-4xMI300) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_USE_DEEP_GEMM=0
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
+
 #---------------------------------------------------------  mi300 · examples  ----------------------------------------------------------#
 
 - label: Examples # TBD
@@ -1419,11 +1547,12 @@ steps:
   commands:
   - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-- label: Kernels Core Operation Test # TBD
+- label: Kernels Core Operation Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  parallelism: 3
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/
@@ -1434,7 +1563,7 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
+  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: Kernels KDA Test # TBD
   timeout_in_minutes: 180
@@ -1452,6 +1581,21 @@ steps:
   commands:
   - pytest -v -s kernels/test_kda.py
 
+- label: Kernels Mamba Test # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/mamba/
+  - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/ops
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/mamba
+
 - label: Kernels MoE Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1593,7 +1737,7 @@ steps:
   - set -x
   - export VLLM_USE_V2_MODEL_RUNNER=1
   - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
   - pytest -v -s v1/e2e/general/test_context_length.py
   - pytest -v -s v1/e2e/general/test_min_tokens.py
   - 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"
@@ -1797,6 +1941,37 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+- label: Language Models Test (Extended Generation) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - 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)'
+
+- label: Language Models Tests (Hybrid) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  parallelism: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
 
 - label: Multi-Modal Models (Extended Generation 1) # TBD
@@ -1899,20 +2074,32 @@ steps:
   commands:
   - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
-- label: Multi-Modal Processor (CPU) %N # TBD
+- label: Multi-Modal Models (Extended Pooling) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/multimodal/pooling
+  commands:
+  - pytest -v -s models/multimodal/pooling -m 'not core_model'
+
+- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
-  parallelism: 4
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
-  - tests/models/registry.py
   commands:
-  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
+  - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
 
 #-----------------------------------------------------  mi300 · models / quantized  -----------------------------------------------------#
 
@@ -1933,12 +2120,12 @@ steps:
 
 #--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#
 
-- label: Transformers Nightly Models (Shardable) %N # TBD
+- label: Transformers Nightly Models (Initialization) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
-  parallelism: 4
+  parallelism: 6
   optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
@@ -1954,6 +2141,27 @@ steps:
   commands:
   - pip install --upgrade git+https://github.com/huggingface/transformers
   - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: Transformers Nightly Models (Processing) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  parallelism: 8
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/multimodal/
+  - vllm/model_executor/layers/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - tests/models/
+  commands:
+  - pip install --upgrade git+https://github.com/huggingface/transformers
   - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: Transformers Nightly Models (Single) # TBD
@@ -2472,11 +2680,12 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/kv_connector/extract_hidden_states_integration
 
-- label: V1 attention (H100-MI300) # TBD
+- label: V1 attention (H100-MI300) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
+  parallelism: 2
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -2488,7 +2697,7 @@ steps:
   - vllm/envs.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s v1/attention
+  - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: V1 Core + KV + Metrics # TBD
   timeout_in_minutes: 180
@@ -2517,23 +2726,6 @@ steps:
   # - export HSA_NO_SCRATCH_RECLAIM=1
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
-- label: V1 others (CPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/v1
-  commands:
-  - 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 -m 'cpu_test' v1/kv_connector/unit
-  - pytest -v -s -m 'cpu_test' v1/metrics
-
 - label: V1 Sample + Logits # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2830,195 +3022,6 @@ steps:
   commands:
   - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
 
-#########################################################################################################################################
-#                                                                                                                                       #
-#                                                         MI325 (gfx942) tests                                                          #
-#                                                                                                                                       #
-#########################################################################################################################################
-
-#----------------------------------------------------------  mi325 · compile  ----------------------------------------------------------#
-
-- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/compile/test_wrapper.py
-  - tests/entrypoints/llm/test_collective_rpc.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
-
-#--------------------------------------------------------  mi325 · distributed  --------------------------------------------------------#
-
-- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/distributed/
-  - tests/v1/shutdown
-  - tests/v1/worker/test_worker_memory_snapshot.py
-  - vllm/platforms/rocm.py
-  commands:
-  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
-  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
-
-- label: Distributed Compile + Comm (4 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - tests/distributed/test_pynccl
-  - tests/distributed/test_events
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/distributed/test_symm_mem_allreduce.py
-  - tests/distributed/test_multiproc_executor.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s distributed/test_pynccl.py
-  - pytest -v -s distributed/test_events.py
-  - pytest -v -s distributed/test_symm_mem_allreduce.py
-  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
-
-#----------------------------------------------------------  mi325 · engine  -----------------------------------------------------------#
-
-- label: Engine # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/engine
-  - tests/test_sequence
-  - tests/test_config
-  - tests/test_logger
-  - tests/test_vllm_port
-  commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
-
-#-----------------------------------------------------------  mi325 · evals  -----------------------------------------------------------#
-
-- label: LM Eval Large Models (4xH100-4xMI325) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_USE_DEEP_GEMM=0
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
-
-#-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#
-
-- label: Language Models Test (Extended Generation) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-  - 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)'
-
-- label: Language Models Tests (Hybrid) %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  parallelism: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
-#----------------------------------------------------  mi325 · models / multimodal  ----------------------------------------------------#
-
-- label: Multi-Modal Models (Extended Pooling) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal/pooling
-  commands:
-  - pytest -v -s models/multimodal/pooling -m 'not core_model'
-
-- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
-  - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
-
-#-----------------------------------------------------------  mi325 · misc  ------------------------------------------------------------#
-
-- label: Python-only Installation # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - tests/standalone_tests/python_only_compile.sh
-  - setup.py
-  - vllm/platforms/rocm.py
-  commands:
-  - bash standalone_tests/python_only_compile.sh
-
 #########################################################################################################################################
 #                                                                                                                                       #
 #                                                         MI355 (gfx950) tests                                                          #
@@ -3030,6 +3033,7 @@ steps:
 - label: Attention Benchmarks Smoke Test (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   working_dir: "/vllm-workspace/"
@@ -3046,6 +3050,7 @@ steps:
 - label: Distributed Tests (2xH100-2xMI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   optional: true
@@ -3090,6 +3095,7 @@ steps:
 - label: Entrypoints Integration (API Server) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   fast_check: true
@@ -3107,6 +3113,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   optional: true
@@ -3122,6 +3129,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   optional: true
@@ -3138,6 +3146,7 @@ steps:
 - label: Entrypoints Integration (API Server Generate) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   optional: true
@@ -3158,6 +3167,7 @@ steps:
 - label: Entrypoints Integration (Speech to Text) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   working_dir: "/vllm-workspace/tests"
@@ -3171,6 +3181,7 @@ steps:
 - label: Entrypoints Integration (Multimodal)
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   working_dir: "/vllm-workspace/tests"
@@ -3184,6 +3195,7 @@ steps:
 - label: Entrypoints Integration (Pooling) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   fast_check: true
   working_dir: "/vllm-workspace/tests"
@@ -3199,6 +3211,7 @@ steps:
 - label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   optional: true
@@ -3221,6 +3234,7 @@ steps:
 - label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   optional: true
@@ -3243,6 +3257,7 @@ steps:
 - label: LM Eval Small Models (2xB200-2xMI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   optional: true
@@ -3262,6 +3277,7 @@ steps:
 - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   working_dir: "/vllm-workspace"
@@ -3282,6 +3298,7 @@ steps:
 - label: LM Eval Large Models (4xH100-4xMI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_4
   num_gpus: 4
   optional: true
@@ -3304,6 +3321,7 @@ steps:
 - label: Examples # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/examples"
   source_file_dependencies:
@@ -3339,6 +3357,7 @@ steps:
 - label: Kernels (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
@@ -3364,6 +3383,7 @@ steps:
 - label: Kernels Attention Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
@@ -3381,6 +3401,7 @@ steps:
 - label: Kernels MoE Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   parallelism: 5
   working_dir: "/vllm-workspace/tests"
@@ -3401,6 +3422,7 @@ steps:
 - label: Kernels Quantization Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
@@ -3418,6 +3440,7 @@ steps:
 - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   working_dir: "/vllm-workspace/tests"
@@ -3437,6 +3460,7 @@ steps:
 - label: Language Models Test (Extended Generation) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3450,6 +3474,7 @@ steps:
 - label: Language Models Test (Extended Pooling) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3462,7 +3487,9 @@ steps:
 - label: Language Models Test (PPL) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/model_executor/models/qwen3_5.py
@@ -3489,6 +3516,7 @@ steps:
 - label: Language Models Tests (Standard) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3503,6 +3531,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3517,6 +3546,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 3) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3529,6 +3559,7 @@ steps:
 - label: Multi-Modal Models (Extended Pooling) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3541,6 +3572,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3554,6 +3586,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3570,6 +3603,7 @@ steps:
 - label: Quantized Models Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3586,6 +3620,7 @@ steps:
 - label: Quantization # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3602,6 +3637,7 @@ steps:
 # - label: Quantized MoE Test (B200-MI355) # TBD
 #   timeout_in_minutes: 180
 #   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+#   dind: false
 #   agent_pool: mi355_1
 #   working_dir: "/vllm-workspace/"
 #   source_file_dependencies:
@@ -3627,10 +3663,12 @@ steps:
 
 #------------------------------------------------------------  mi355 · v1  -------------------------------------------------------------#
 
-- label: V1 attention (B200-MI355) # TBD
+- label: V1 attention (B200-MI355) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
+  parallelism: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/config/attention.py
@@ -3641,11 +3679,12 @@ steps:
   - vllm/envs.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s v1/attention
+  - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: V1 Core + KV + Metrics # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3672,6 +3711,7 @@ steps:
 - label: V1 Sample + Logits # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -3692,6 +3732,7 @@ steps:
 - label: V1 Spec Decode # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3705,6 +3746,7 @@ steps:
 - label: Weight Loading Multiple GPU # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   num_gpus: 2
   working_dir: "/vllm-workspace/tests"
@@ -3717,6 +3759,7 @@ steps:
 - label: Weight Loading Multiple GPU - Large Models # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_2
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
@@ -3732,6 +3775,7 @@ steps:
 - label: Regression # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
   agent_pool: mi355_1
   optional: true
   working_dir: "/vllm-workspace/tests"

.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,8 +16,10 @@ steps:
   commands:
     - pytest -v -s cuda/test_cuda_context.py
     - pytest -v -s cuda/test_platform_no_cuda_init.py
+    - pytest -v -s cuda/test_cuda_compatibility_path.py
 
 - label: Cudagraph
+  device: h200_35gb
   key: cudagraph
   timeout_in_minutes: 30
   source_file_dependencies:
@@ -25,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_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

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: Entrypoints Unit Tests
+  device: h200_35gb
   key: entrypoints-unit-tests
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/tests"
@@ -15,6 +16,7 @@ steps:
   - pytest -v -s entrypoints/weight_transfer
 
 - label: Entrypoints Integration (LLM)
+  device: h200_35gb
   key: entrypoints-integration-llm
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"
@@ -28,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/
@@ -50,13 +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/
@@ -67,14 +72,16 @@ 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
 
 - 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/
@@ -86,12 +93,14 @@ 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
 
 - label: Entrypoints Integration (API Server Generate)
+  device: h200_35gb
   key: entrypoints-integration-api-server-generate
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
@@ -108,12 +117,14 @@ 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
 
 - label: Entrypoints Integration (Responses API)
+  device: h200_35gb
   key: entrypoints-integration-responses-api
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
@@ -148,8 +159,9 @@ steps:
   - pytest -v -s entrypoints/multimodal
 
 - 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/
@@ -169,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

@@ -15,6 +15,7 @@ steps:
     - pytest -v -s tests/kernels/ir
 
 - label: Kernels Core Operation Test
+  device: h200_35gb
   key: kernels-core-operation-test
   timeout_in_minutes: 120
   source_file_dependencies:
@@ -60,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
@@ -79,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
@@ -117,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
@@ -147,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/
@@ -163,6 +204,7 @@ steps:
       - image-build-amd
 
 - label: Kernels Mamba Test
+  device: h200_35gb
   key: kernels-mamba-test
   timeout_in_minutes: 40
   source_file_dependencies:
@@ -235,6 +277,11 @@ steps:
   - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
   - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
   - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
+  - vllm/cute_utils/
+  - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
+  - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
+  - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
+  - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
   - tests/kernels/test_ll_bf16_gemm.py
   - tests/kernels/test_top_k_per_row.py
   commands:
@@ -264,6 +311,8 @@ steps:
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
     - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
+    - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
+    - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
     - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
     # e2e
     - pytest -v -s tests/models/quantization/test_nvfp4.py

.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:
@@ -78,6 +79,28 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
 
+- label: LM Eval PCP (4xB200)
+  key: lm-eval-pcp-4xb200
+  timeout_in_minutes: 360
+  device: b200-k8s
+  num_devices: 4
+  optional: true
+  source_file_dependencies:
+  - csrc/
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+  - tests/evals/gsm8k/configs/models-pcp.txt
+  - vllm/model_executor/layers/quantization
+  - vllm/config/parallel.py
+  - vllm/distributed/parallel_state.py
+  - vllm/model_executor/layers/attention/mla_attention.py
+  - vllm/model_executor/layers/attention/pcp.py
+  - vllm/v1/worker/gpu/model_runner.py
+  - vllm/v1/worker/gpu/pcp_manager.py
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
+
 - label: LM Eval Large Models EP (2xB200)
   key: lm-eval-large-models-ep-2xb200
   timeout_in_minutes: 60
@@ -119,7 +142,7 @@ steps:
     amd:
       dind: false
       device: mi300_8
-      timeout_in_minutes: 60
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
       commands:
@@ -314,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:
@@ -324,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,11 @@ steps:
   parallelism: 4
   mirror:
     amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      soft_fail: true
       working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 65
+      timeout_in_minutes: 85
       source_file_dependencies:
       - vllm/lora
       - tests/lora
@@ -46,4 +48,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,13 +59,16 @@ 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
 
 - label: V1 Core + KV + Metrics
+  device: h200_35gb
   key: v1-core-kv-metrics
-  timeout_in_minutes: 60
+  timeout_in_minutes: 80
   source_file_dependencies:
     - vllm/config/
     - vllm/distributed/
@@ -89,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
@@ -102,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
@@ -110,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
 
@@ -142,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
 
@@ -205,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
@@ -229,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
@@ -256,6 +266,7 @@ steps:
   - vllm/utils/
   - vllm/v1/
   - tests/v1/tracing
+  - tests/tracing/
   commands:
   - "pip install \
       'opentelemetry-sdk>=1.26.0' \
@@ -263,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
@@ -283,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:
@@ -384,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
@@ -394,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
@@ -409,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
@@ -425,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

@@ -3,13 +3,16 @@ depends_on:
   - image-build
 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
@@ -27,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,37 @@ 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/
+  commands:
+    # The native NVIDIA AttnRes kernel requires the SM100 family.
+    - pytest -v -s models/kimi_k3
+
 - label: Basic Models Test (Other CPU) # 5min
   key: basic-models-test-other-cpu
   depends_on:
@@ -55,7 +82,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,10 +17,12 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 45
       depends_on:
       - image-build-amd
 
 - label: Language Models Tests (Extra Standard) %N
+  device: h200_35gb
   key: language-models-tests-extra-standard
   timeout_in_minutes: 40
   source_file_dependencies:
@@ -38,6 +40,7 @@ steps:
     amd:
       dind: false
       device: mi300_1
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -51,8 +54,8 @@ steps:
       - tests/models/language/pooling/test_classification.py
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
-
 - label: Language Models Tests (Hybrid) %N
+  device: h200_35gb
   key: language-models-tests-hybrid
   timeout_in_minutes: 65
   source_file_dependencies:
@@ -60,16 +63,16 @@ steps:
   - tests/models/language/generation
   commands:
     # Install fast path packages for testing against transformers
-    # Note: also needed to run plamo2 model in vLLM
     - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    # Shard hybrid language model tests
-    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+    # Shard the hybrid language model tests that are numerically stable on Hopper.
+    - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   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:
@@ -77,6 +80,20 @@ steps:
       - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
       - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
+# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
+# correctness test on L4 until its H200 output matches the Transformers reference.
+- label: Language Models Tests (Granite L4 Compatibility)
+  key: language-models-tests-granite-l4-compatibility
+  timeout_in_minutes: 65
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+    - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
+
 - label: Language Models Test (Extended Generation) # 80min
   device: h200_35gb
   key: language-models-test-extended-generation
@@ -87,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)'
@@ -115,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:
@@ -119,6 +127,7 @@ steps:
       - vllm/model_executor/model_loader/
 
 - label: Multi-Modal Models (Extended Generation 1)
+  device: h200_35gb
   key: multi-modal-models-extended-generation-1
   optional: true
   source_file_dependencies:
@@ -130,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
 
@@ -165,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,17 +107,11 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
-  mirror:
-    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd
 
 - label: PyTorch Fullgraph Smoke Test
+  device: h200_35gb
   key: pytorch-fullgraph-smoke-test
-  timeout_in_minutes: 60
+  timeout_in_minutes: 90
   source_file_dependencies:
   - vllm/__init__.py
   - vllm/_aiter_ops.py
@@ -149,7 +143,42 @@ steps:
   # as it is a heavy test that is covered in other steps.
   # Use `find` to launch multiple instances of pytest so that
   # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+
+# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
+# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
+# those SM90 cases skip while the architecture-compatible cases still run.
+- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
+  key: pytorch-fullgraph-cudagraph-l4-compatibility
+  timeout_in_minutes: 60
+  source_file_dependencies:
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/compile
+  commands:
+  - pytest -s -v compile/fullgraph/test_full_cudagraph.py
 
 - label: PyTorch Fullgraph
   key: pytorch-fullgraph
@@ -200,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

@@ -3,8 +3,11 @@ depends_on:
   - image-build
 steps:
 - label: Quantization
+  device: h200_35gb
   key: quantization
-  timeout_in_minutes: 60
+  timeout_in_minutes: 75
+  env:
+    VLLM_USE_V2_MODEL_RUNNER: "0"
   source_file_dependencies:
   - csrc/
   - vllm/model_executor/layers/quantization
@@ -19,9 +22,12 @@ steps:
   # TODO(jerryzh168): resolve the above comment
   - uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
   - uv pip install --system conch-triton-kernels
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
+  # 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'
 
 - label: Quantized Fusions
+  device: h200_35gb
   key: quantized-fusions
   timeout_in_minutes: 20
   source_file_dependencies:
@@ -52,8 +58,11 @@ steps:
     - pytest -s -v tests/quantization/test_blackwell_moe.py
 
 - label: Quantized Models Test
+  device: h200_35gb
   key: quantized-models-test
-  timeout_in_minutes: 50
+  timeout_in_minutes: 65
+  env:
+    VLLM_USE_V2_MODEL_RUNNER: "0"
   source_file_dependencies:
   - vllm/model_executor/layers/quantization
   - tests/models/quantization

.buildkite/test_areas/rust_frontend.yaml

@@ -81,6 +81,7 @@ steps:
   - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 - label: Rust Frontend Tool Use
+  device: h200_35gb
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:

.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(', ')}`);

.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

@@ -30,7 +30,7 @@ repos:
   - id: markdownlint-cli2
     language_version: lts
     args: [--fix]
-    exclude: ^CLAUDE\.md$
+    exclude: (^|/)CLAUDE\.md$
 - repo: https://github.com/rhysd/actionlint
   rev: v1.7.7
   hooks:
@@ -260,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
@@ -420,7 +430,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 +705,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 +825,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 +909,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 +934,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 +991,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 +1057,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 +1079,24 @@ 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(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 +1138,10 @@ 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(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)

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/bench_cp_gather_fp8.py

@@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size):
     # Output workspace
     dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
 
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
     workspace_starts_t = torch.tensor(
         workspace_starts, dtype=torch.int32, device="cuda"
     )
 
-    return cache, dst, block_table, seq_lens_t, workspace_starts_t
+    return cache, dst, block_table, workspace_starts_t
 
 
 def bench_scenario(label, num_reqs, total_tokens_list, save_path):
@@ -94,7 +93,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
         )
     )
     def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
+        cache, dst, block_table, ws_starts = make_inputs(
             total_tokens, num_reqs, BLOCK_SIZE
         )
 
@@ -102,7 +101,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
 
         ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
             lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
+                cache, dst, block_table, ws_starts, num_reqs
             ),
             quantiles=quantiles,
             rep=500,

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

@@ -14,7 +14,7 @@ else()
   FetchContent_Declare(
     tml_fa4
     GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
-    GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e
+    GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND "")

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/cache.h

@@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
     torch::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& seq_lens,          // [BATCH]
     torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
+    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
 
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(

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_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/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);
 
@@ -207,12 +207,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 +516,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
@@ -570,8 +611,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 +622,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 +650,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/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/cache_kernels.cu

@@ -1025,6 +1025,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 =
@@ -1174,7 +1177,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
     const int32_t num_reqs, const int32_t block_size,
     const int32_t total_tokens, const int64_t block_table_stride,
     const int64_t cache_block_stride, const int64_t cache_entry_stride,
-    const int64_t dst_entry_stride) {
+    const int64_t dst_entry_stride,
+    const int32_t* __restrict__ seq_starts) {  // Optional source offsets
   const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
   if (flat_warp_id >= total_tokens) return;
   const int lane_id = threadIdx.x & 31;
@@ -1192,7 +1196,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
 
   // Compute physical token address via block table
   const int out_token_id = flat_warp_id;
-  const int token_offset = out_token_id - workspace_starts[req_id];
+  int token_offset = out_token_id - workspace_starts[req_id];
+  if (seq_starts != nullptr) token_offset += seq_starts[req_id];
   const int cache_block_idx = token_offset / block_size;
   const int offset_in_block = token_offset % block_size;
   const int physical_block =
@@ -1383,9 +1388,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::stable::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, 656]
     torch::stable::Tensor const& dst,          // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& seq_lens,     // [BATCH]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size) {
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
   torch::stable::accelerator::DeviceGuard device_guard(
       src_cache.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
@@ -1396,20 +1401,25 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   STD_TORCH_CHECK(
       block_table.scalar_type() == torch::headeronly::ScalarType::Int,
       "block_table must be int32");
-  STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
-                  "seq_lens must be int32");
   STD_TORCH_CHECK(
       workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
       "workspace_starts must be int32");
+  if (seq_starts.has_value()) {
+    STD_TORCH_CHECK(
+        seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
+        "seq_starts must be int32");
+  }
 
   STD_TORCH_CHECK(src_cache.device() == dst.device(),
                   "src_cache and dst must be on the same device");
   STD_TORCH_CHECK(src_cache.device() == block_table.device(),
                   "src_cache and block_table must be on the same device");
-  STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),
-                  "src_cache and seq_lens must be on the same device");
   STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
                   "src_cache and workspace_starts must be on the same device");
+  if (seq_starts.has_value()) {
+    STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
+                    "src_cache and seq_starts must be on the same device");
+  }
   auto dtype = src_cache.scalar_type();
   STD_TORCH_CHECK(
       dtype == torch::headeronly::ScalarType::Byte ||               // uint8
@@ -1438,6 +1448,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   constexpr int warps_per_block = 8;
   const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
   const int block_size_threads = warps_per_block * 32;  // 256 threads
+  const int32_t* seq_starts_ptr =
+      seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
+                             : nullptr;
 
   vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
                                                stream>>>(
@@ -1446,7 +1459,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
       workspace_starts.const_data_ptr<int32_t>(),
       static_cast<int32_t>(batch_size), block_size, total_tokens,
       block_table_stride, cache_block_stride, cache_entry_stride,
-      dst_entry_stride);
+      dst_entry_stride, seq_starts_ptr);
 }
 
 // Macro to dispatch the kernel based on the data type.

csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu

@@ -0,0 +1,954 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ */
+
+// Production AttnRes forward for Blackwell (SM100).
+//
+// Warp-specialized online softmax + residual + RMSNorm:
+//   - 1 producer warp issues cp.async.bulk row loads into shared memory.
+//   - 8 consumer warps compute reductions and output.
+//   - Q=res_weight*rms_weight remains in registers across persistent tokens.
+//   - V rows are converted once and cached as FP32 in TMEM between passes.
+//
+// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major
+// block residual storage.
+
+#include "../torch_utils.h"
+
+#include <cfloat>
+#include <cstdint>
+#include <cstdio>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+using bf16_t = __nv_bfloat16;
+
+namespace sm100 {
+namespace fwd_prod_v2 {
+
+constexpr int K_TILE = 1024;
+constexpr int N_CHUNK_DEFAULT = 4;
+constexpr int CHUNK_DEPTH = 2;
+constexpr int BLK = 288;  // 1 producer warp + 8 consumer warps
+constexpr int CONSUMER_THREADS = BLK - 32;  // 256
+constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32;
+constexpr int CONSUMER_GROUPS = 2;  // two 128-thread consumer groups
+constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS;
+constexpr int FIRST_USER_NAMED_BARRIER = 8;
+
+__device__ __forceinline__ const bf16_t* residual_addr(
+    const bf16_t* block_res, const bf16_t* layer_res, int source, int N,
+    int token, int block_stride_m, int block_stride_r, int H) {
+  if (source < N - 1) {
+    return block_res + static_cast<long long>(token) * block_stride_m +
+           source * block_stride_r;
+  }
+  return layer_res + static_cast<long long>(token) * H;
+}
+
+__device__ __forceinline__ uint32_t elect_one_sync() {
+  uint32_t pred = 0;
+  uint32_t laneid = 0;
+  asm volatile(
+      "{\n"
+      ".reg .b32 %%rx;\n"
+      ".reg .pred %%px;\n"
+      "     elect.sync %%rx|%%px, %2;\n"
+      "@%%px mov.s32 %1, 1;\n"
+      "     mov.s32 %0, %%rx;\n"
+      "}\n"
+      : "+r"(laneid), "+r"(pred)
+      : "r"(0xffffffff));
+  return pred;
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t& barrier,
+                                              int thread_count) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr),
+               "r"(thread_count));
+}
+
+__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier,
+                                                   uint32_t bytes) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(
+                   barrier_addr),
+               "r"(bytes));
+}
+
+__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile(
+      "{\n"
+      ".reg .pred p;\n"
+      "WAIT:\n"
+      "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
+      "@p bra DONE;\n"
+      "bra WAIT;\n"
+      "DONE:\n"
+      "}\n" ::"r"(barrier_addr),
+      "r"(phase));
+}
+
+__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile(
+      "{\n"
+      ".reg .b64 state;\n"
+      "mbarrier.arrive.shared::cta.b64 state, [%0];\n"
+      "}\n" ::"r"(barrier_addr));
+}
+
+__device__ __forceinline__ void fence_mbarrier_init() {
+  asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+}
+
+__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads,
+                                                   uint32_t user_barrier_id) {
+  asm volatile(
+      "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER),
+      "r"(num_threads)
+      : "memory");
+}
+
+__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) {
+  uint32_t const dst_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(dst));
+  asm volatile(
+      "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
+          dst_addr),
+      "r"(num_columns));
+}
+
+__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) {
+  asm volatile(
+      "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr),
+      "r"(num_columns));
+}
+
+__device__ __forceinline__ void tmem_release_allocation_lock() {
+  asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
+}
+
+__device__ __forceinline__ void tmem_store_wait() {
+  asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
+}
+
+template <int N, typename T>
+__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) {
+  uint32_t* values = reinterpret_cast<uint32_t*>(dst);
+  if constexpr (N == 8) {
+    asm volatile(
+        "tcgen05.ld.sync.aligned.32x32b.x8.b32"
+        "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
+        : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]),
+          "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7])
+        : "r"(src_addr));
+  } else {
+    static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
+    asm volatile(
+        "tcgen05.ld.sync.aligned.32x32b.x4.b32"
+        "{%0, %1, %2, %3}, [%4];\n"
+        : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3])
+        : "r"(src_addr));
+  }
+}
+
+template <int N, typename T>
+__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) {
+  uint32_t* values = reinterpret_cast<uint32_t*>(src);
+  if constexpr (N == 8) {
+    asm volatile(
+        "tcgen05.st.sync.aligned.32x32b.x8.b32"
+        "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]),
+        "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]),
+        "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr));
+  } else {
+    static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
+    asm volatile(
+        "tcgen05.st.sync.aligned.32x32b.x4.b32"
+        "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]),
+        "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr));
+  }
+}
+
+__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) {
+  float2 result;
+  asm volatile("add.rn.f32x2 %0, %1, %2;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)));
+  return result;
+}
+
+__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) {
+  float2 result;
+  asm volatile("mul.f32x2 %0, %1, %2;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)));
+  return result;
+}
+
+__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b,
+                                             const float2& c) {
+  float2 result;
+  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)),
+                 "l"(reinterpret_cast<uint64_t const&>(c)));
+  return result;
+}
+
+template <int NC>
+struct FwdSmemPlan {
+  alignas(16) uint64_t bar_ready[CHUNK_DEPTH];
+  alignas(16) uint64_t bar_consumed[CHUNK_DEPTH];
+  alignas(16) uint64_t bar_output_norm_ready;
+  alignas(16) float2 ws_stats[CONSUMER_WARPS][NC];
+  uint32_t tmem_base;
+};
+
+__device__ __forceinline__ void cp_async_bulk(void* smem_dst,
+                                              const void* gmem_src, int bytes,
+                                              uint64_t& mbar) {
+  uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));
+  uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar));
+  asm volatile(
+      "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], "
+      "[%1], %2, [%3];\n" ::"r"(s),
+      "l"(gmem_src), "r"(bytes), "r"(m)
+      : "memory");
+}
+
+template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
+          bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
+          bool OUTPUT_NORM_IN_SMEM = false>
+__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
+    const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
+    const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
+    const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
+    int B, int block_stride_m, int block_stride_r, float rms_eps,
+    const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
+  constexpr float LOG2_E = 1.4426950408889634f;
+  constexpr int N_CHUNK = NC;
+  // The two-source specialization only consumes half of the TMEM columns.
+  constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256;
+  constexpr int NUM_BUFS = CHUNK_DEPTH * NC;
+  constexpr int NHT = H / K_TILE;
+  constexpr int SLICES_PER_GROUP =
+      (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS;
+  constexpr int VEC = 8;
+  constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC;
+  constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC;
+  constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP;
+  static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
+  static_assert(H >= 4096 && H <= 8192);
+  static_assert(H % K_TILE == 0);
+
+  const int tid = threadIdx.x;
+  const int wid = tid >> 5;
+  const int lane = tid & 31;
+  const int TB = T * B;
+  const int num_ctas = gridDim.x;
+  const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
+
+  const int comp_wid = wid - 1;
+  const int comp_tid = tid - 32;
+  const int group = (comp_wid >= 4) ? 1 : 0;
+  const int ct_in_group =
+      (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1;
+  const int k_local = ct_in_group * VEC;
+
+  constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
+  constexpr size_t DELTA_BYTES =
+      HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
+  constexpr size_t OUTPUT_NORM_BYTES =
+      OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
+  extern __shared__ __align__(16) char smem_raw[];
+  bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw);  // [NUM_BUFS][H]
+  bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
+  bf16_t* output_norm_buf =
+      reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES);
+  FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>(
+      smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES);
+
+  auto slot_of = [](long long gci, int n) {
+    return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n;
+  };
+  auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); };
+  auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; };
+  auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* {
+    return delta_bufs + chunk_slot * H;
+  };
+
+  if (wid == 0 && elect_one_sync()) {
+  #pragma unroll
+    for (int i = 0; i < CHUNK_DEPTH; i++) {
+      mbarrier_init(plan.bar_ready[i], 1);
+      mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
+    }
+    if constexpr (OUTPUT_NORM_IN_SMEM) {
+      mbarrier_init(plan.bar_output_norm_ready, 1);
+    }
+    fence_mbarrier_init();
+  }
+
+  // gdc wait BEFORE tmem alloc
+  cudaGridDependencySynchronize();
+
+  if (wid == 1) {
+    tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
+    if constexpr (RELEASE_TMEM) {
+      tmem_release_allocation_lock();
+    }
+  }
+  __syncthreads();
+
+  if constexpr (OUTPUT_NORM_IN_SMEM) {
+    if (wid == 0 && elect_one_sync()) {
+      mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
+      cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
+                    plan.bar_output_norm_ready);
+    }
+  }
+
+  const uint32_t my_v_tmem =
+      comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0;
+  float q_cache[ACC_PER_THREAD];
+  if (comp_tid >= 0) {
+  #pragma unroll
+    for (int si = 0; si < SLICES_PER_GROUP; si++) {
+      if constexpr (H == 7168) {
+        if (si == SLICES_PER_GROUP - 1) {
+          int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+  #pragma unroll
+          for (int j = 0; j < 4; j++) {
+            int h = h_base + j;
+            q_cache[si * VEC + j] =
+                __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
+          }
+          continue;
+        }
+      }
+      int dt = si * CONSUMER_GROUPS + group;
+      if (dt >= NHT) continue;
+      int h_base = dt * K_TILE + k_local;
+  #pragma unroll
+      for (int j = 0; j < VEC; j++) {
+        int h = h_base + j;
+        q_cache[si * VEC + j] =
+            __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
+      }
+    }
+  }
+
+  if (wid == 0) {
+    if (elect_one_sync()) {
+      long long gci = 0;
+      for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
+        const int t = tb / B;
+        for (int ci = 0; ci < num_chunks; ci++, gci++) {
+          int ns = ci * N_CHUNK;
+          int an = min(N_CHUNK, N - ns);
+          int chunk_slot = (int)(gci % CHUNK_DEPTH);
+          int pc = phase_of(gci);
+          mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
+          int transaction_bytes = an * H * (int)sizeof(bf16_t);
+          if constexpr (HAS_DELTA) {
+            int prefix_n = N - 1 - ns;
+            if (prefix_n >= 0 && prefix_n < an) {
+              transaction_bytes += H * (int)sizeof(bf16_t);
+            }
+          }
+          mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes);
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            if (n >= an) continue;
+            int slot = slot_of(gci, n);
+            const bf16_t* src =
+                residual_addr(block_res, layer_res, ns + n, N, t,
+                              block_stride_m, block_stride_r, H);
+            cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
+                          plan.bar_ready[chunk_slot]);
+          }
+          if constexpr (HAS_DELTA) {
+            int prefix_n = N - 1 - ns;
+            if (prefix_n >= 0 && prefix_n < an) {
+              cp_async_bulk(delta_buf_ptr(chunk_slot),
+                            delta + (long long)tb * H, H * sizeof(bf16_t),
+                            plan.bar_ready[chunk_slot]);
+            }
+          }
+        }
+      }
+    }
+  } else {
+    float acc32[ACC_PER_THREAD] = {};
+    float eps_cache;
+    asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps));
+
+    long long gci = 0;
+    for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
+      float m_running = -FLT_MAX;
+      float s_running = 0.f;
+  #pragma unroll
+      for (int i = 0; i < ACC_PER_THREAD; i++) {
+        acc32[i] = 0.f;
+      }
+
+      for (int ci = 0; ci < num_chunks; ci++, gci++) {
+        int ns = ci * N_CHUNK;
+        int an = min(N_CHUNK, N - ns);
+        int chunk_slot = (int)(gci % CHUNK_DEPTH);
+        int pr = phase_of(gci);
+        mbarrier_wait(plan.bar_ready[chunk_slot], pr);
+
+        float2 sq_local[N_CHUNK] = {};
+        float2 dot_local[N_CHUNK] = {};
+
+        auto pass_A_body = [&](auto AN_TOK) {
+          constexpr int AN = decltype(AN_TOK)::value;
+  #pragma unroll
+          for (int si = 0; si < SLICES_PER_GROUP; si++) {
+            if constexpr (H == 7168) {
+              if (si == SLICES_PER_GROUP - 1) {
+                int h_base =
+                    6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+                const float* qv = &q_cache[si * VEC];
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  int slot = slot_of(gci, n);
+                  int2 vp =
+                      *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
+                  auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
+                  if constexpr (HAS_DELTA) {
+                    int prefix_n = N - 1 - ns;
+                    if (n == prefix_n) {
+                      const bf16_t* delta_ptr =
+                          delta_buf_ptr(chunk_slot) + h_base;
+  #pragma unroll
+                      for (int j = 0; j < 2; j++) {
+                        auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
+                            delta_ptr + 2 * j);
+                        v2[j] = __hadd2(v2[j], delta2);
+                      }
+                      *reinterpret_cast<int2*>(layer_res + (long long)tb * H +
+                                               h_base) = vp;
+                    }
+                  }
+                  float2 f[2] = {__bfloat1622float2(v2[0]),
+                                 __bfloat1622float2(v2[1])};
+                  tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
+                  sq_local[n] = float2_fma(f[0], f[0], sq_local[n]);
+                  sq_local[n] = float2_fma(f[1], f[1], sq_local[n]);
+                  dot_local[n] =
+                      float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]);
+                  dot_local[n] =
+                      float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]);
+                }
+                continue;
+              }
+            }
+            int dt = si * CONSUMER_GROUPS + group;
+            if (dt >= NHT) continue;
+            int h_base = dt * K_TILE + k_local;
+            const float* qv = &q_cache[si * VEC];
+
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              int slot = slot_of(gci, n);
+              int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
+              auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
+              if constexpr (HAS_DELTA) {
+                int prefix_n = N - 1 - ns;
+                if (n == prefix_n) {
+                  const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
+  #pragma unroll
+                  for (int j = 0; j < VEC / 2; j++) {
+                    auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
+                        delta_ptr + 2 * j);
+                    v2[j] = __hadd2(v2[j], delta2);
+                  }
+                  *reinterpret_cast<int4*>(layer_res + (long long)tb * H +
+                                           h_base) = vp;
+                }
+              }
+              float2 f[4] = {
+                  __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
+                  __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])};
+              tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
+  #pragma unroll
+              for (int j = 0; j < VEC / 2; j++) {
+                sq_local[n] = float2_fma(f[j], f[j], sq_local[n]);
+                dot_local[n] = float2_fma(
+                    f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]);
+              }
+            }
+          }
+        };
+        if constexpr (NC == 4) {
+          switch (an) {
+            case 4:
+              pass_A_body(std::integral_constant<int, 4>{});
+              break;
+            case 3:
+              pass_A_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else if constexpr (NC == 3) {
+          switch (an) {
+            case 3:
+              pass_A_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else {
+          static_assert(NC == 2);
+          switch (an) {
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        }
+        if (lane == 0) {
+          mbarrier_arrive(plan.bar_consumed[chunk_slot]);
+        }
+        tmem_store_wait();
+
+        float2 reduce_pair[N_CHUNK];
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y,
+                                       dot_local[n].x + dot_local[n].y);
+        }
+  #pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1) {
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]);
+            packed = __shfl_xor_sync(0xffffffff, packed, offset);
+            float2 other = reinterpret_cast<float2&>(packed);
+            reduce_pair[n] = float2_add(reduce_pair[n], other);
+          }
+        }
+        if (lane == 0) {
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            plan.ws_stats[comp_wid][n] = reduce_pair[n];
+          }
+        }
+        named_barrier_sync(CONSUMER_THREADS, 0);
+
+        float local_rsig = 0.f;
+        float local_logit = 0.f;
+        int stat_n = lane / CONSUMER_WARPS;
+        int stat_w = lane % CONSUMER_WARPS;
+        float2 totals = {};
+        if (stat_n < N_CHUNK) {
+          totals = plan.ws_stats[stat_w][stat_n];
+        }
+  #pragma unroll
+        for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
+          totals.x +=
+              __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS);
+          totals.y +=
+              __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS);
+        }
+        if (stat_n < N_CHUNK && stat_w == 0) {
+          local_rsig = rsqrtf(totals.x / H + eps_cache);
+          local_logit = totals.y * local_rsig;
+        }
+        float logit_n[N_CHUNK];
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS);
+        }
+
+        float m_chunk = -FLT_MAX;
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]);
+        }
+        float m_new = fmaxf(m_running, m_chunk);
+        float corr = exp2f((m_running - m_new) * LOG2_E);
+        float w_n[N_CHUNK] = {};
+        float w_sum = 0.f;
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          if (n < an) {
+            w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E);
+            w_sum += w_n[n];
+          }
+        }
+
+        auto pass_B_body = [&](auto AN_TOK) {
+          constexpr int AN = decltype(AN_TOK)::value;
+  #pragma unroll
+          for (int si = 0; si < SLICES_PER_GROUP; si++) {
+            if constexpr (H == 7168) {
+              if (si == SLICES_PER_GROUP - 1) {
+                float2 corr2 = make_float2(corr, corr);
+                float2 a[2];
+  #pragma unroll
+                for (int j = 0; j < 2; j++) {
+                  float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                           acc32[si * VEC + 2 * j + 1]);
+                  a[j] = float2_mul(old, corr2);
+                }
+                float2 f_cache[AN][2];
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC,
+                               f_cache[n]);
+                }
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  float2 wn = make_float2(w_n[n], w_n[n]);
+  #pragma unroll
+                  for (int j = 0; j < 2; j++) {
+                    a[j] = float2_fma(wn, f_cache[n][j], a[j]);
+                  }
+                }
+  #pragma unroll
+                for (int j = 0; j < 2; j++) {
+                  acc32[si * VEC + 2 * j] = a[j].x;
+                  acc32[si * VEC + 2 * j + 1] = a[j].y;
+                }
+                continue;
+              }
+            }
+            int dt = si * CONSUMER_GROUPS + group;
+            if (dt >= NHT) continue;
+            float2 corr2 = make_float2(corr, corr);
+            float2 a[VEC / 2];
+  #pragma unroll
+            for (int j = 0; j < VEC / 2; j++) {
+              float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                       acc32[si * VEC + 2 * j + 1]);
+              a[j] = float2_mul(old, corr2);
+            }
+            float2 f_cache[AN][VEC / 2];
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]);
+            }
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              float2 wn = make_float2(w_n[n], w_n[n]);
+  #pragma unroll
+              for (int j = 0; j < VEC / 2; j++) {
+                a[j] = float2_fma(wn, f_cache[n][j], a[j]);
+              }
+            }
+  #pragma unroll
+            for (int j = 0; j < VEC / 2; j++) {
+              acc32[si * VEC + 2 * j] = a[j].x;
+              acc32[si * VEC + 2 * j + 1] = a[j].y;
+            }
+          }
+        };
+        if constexpr (NC == 4) {
+          switch (an) {
+            case 4:
+              pass_B_body(std::integral_constant<int, 4>{});
+              break;
+            case 3:
+              pass_B_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else if constexpr (NC == 3) {
+          switch (an) {
+            case 3:
+              pass_B_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else {
+          static_assert(NC == 2);
+          switch (an) {
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        }
+
+        s_running = s_running * corr + w_sum;
+        m_running = m_new;
+      }
+
+      float inv_s = 1.f / s_running;
+      bf16_t* out_ptr = output + (long long)tb * H;
+      float2 output_sq_pair = {};
+      // When output RMSNorm is fused, the softmax denominator cancels:
+      // (acc / s) * rsqrt(mean((acc / s)^2) + eps)
+      //   = acc * rsqrt(mean(acc^2) + eps * s^2).
+  #pragma unroll
+      for (int si = 0; si < SLICES_PER_GROUP; si++) {
+        if constexpr (H == 7168) {
+          if (si == SLICES_PER_GROUP - 1) {
+            int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+            uint2 packed;
+            auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
+            float2 inv2 = make_float2(inv_s, inv_s);
+  #pragma unroll
+            for (int j = 0; j < 2; j++) {
+              float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                       acc32[si * VEC + 2 * j + 1]);
+              if constexpr (HAS_OUTPUT_NORM) {
+                output_sq_pair = float2_fma(old, old, output_sq_pair);
+              } else {
+                float2 mixed = float2_mul(old, inv2);
+                ov2[j] = __float22bfloat162_rn(mixed);
+              }
+            }
+            if constexpr (!HAS_OUTPUT_NORM) {
+              *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
+            }
+            continue;
+          }
+        }
+        int dt = si * CONSUMER_GROUPS + group;
+        if (dt >= NHT) continue;
+        int h_base = dt * K_TILE + k_local;
+        uint4 packed;
+        auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
+        float2 inv2 = make_float2(inv_s, inv_s);
+  #pragma unroll
+        for (int j = 0; j < VEC / 2; j++) {
+          float2 old =
+              make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]);
+          if constexpr (HAS_OUTPUT_NORM) {
+            output_sq_pair = float2_fma(old, old, output_sq_pair);
+          } else {
+            float2 mixed = float2_mul(old, inv2);
+            ov2[j] = __float22bfloat162_rn(mixed);
+          }
+        }
+        if constexpr (!HAS_OUTPUT_NORM) {
+          *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
+        }
+      }
+
+      if constexpr (HAS_OUTPUT_NORM) {
+        if constexpr (OUTPUT_NORM_IN_SMEM) {
+          // The immutable weight copy is acquired once, at its first use.
+          if (tb == blockIdx.x) {
+            mbarrier_wait(plan.bar_output_norm_ready, 0);
+          }
+        }
+        float output_sq = output_sq_pair.x + output_sq_pair.y;
+  #pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1) {
+          output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset);
+        }
+        if (lane == 0) {
+          plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f);
+        }
+        named_barrier_sync(CONSUMER_THREADS, 0);
+        float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f;
+  #pragma unroll
+        for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
+          total_sq +=
+              __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS);
+        }
+        if (lane == 0) {
+          total_sq =
+              rsqrtf(total_sq / H + output_norm_eps * s_running * s_running);
+        }
+        float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0);
+  #pragma unroll
+        for (int si = 0; si < SLICES_PER_GROUP; si++) {
+          if constexpr (H == 7168) {
+            if (si == SLICES_PER_GROUP - 1) {
+              int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+              uint2 packed;
+              auto* values = reinterpret_cast<bf16_t*>(&packed);
+  #pragma unroll
+              for (int j = 0; j < 4; j++) {
+                const bf16_t* weight_ptr =
+                    OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
+                float weight = __bfloat162float(weight_ptr[h_base + j]);
+                values[j] = __float2bfloat16(acc32[si * VEC + j] *
+                                             output_rsigma * weight);
+              }
+              *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
+              continue;
+            }
+          }
+          int dt = si * CONSUMER_GROUPS + group;
+          if (dt >= NHT) continue;
+          int h_base = dt * K_TILE + k_local;
+          uint4 packed;
+          auto* values = reinterpret_cast<bf16_t*>(&packed);
+  #pragma unroll
+          for (int j = 0; j < VEC; j++) {
+            const bf16_t* weight_ptr =
+                OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
+            float weight = __bfloat162float(weight_ptr[h_base + j]);
+            values[j] =
+                __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
+          }
+          *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
+        }
+      }
+    }
+  }
+
+  cudaTriggerProgrammaticLaunchCompletion();
+  __syncthreads();
+  if (wid == 1) {
+    tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
+  }
+#else
+  if (threadIdx.x == 0) {
+    printf("attn_res_fwd_online_v2_kernel requires sm_10x\n");
+  }
+#endif
+}
+
+template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
+          bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
+          bool OUTPUT_NORM_IN_SMEM = false>
+static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
+                       const bf16_t* delta, const bf16_t* res_weight,
+                       const bf16_t* rms_weight, bf16_t* output, int N, int T,
+                       int B, float rms_eps, int num_sm, cudaStream_t stream,
+                       const bf16_t* output_norm_weight = nullptr,
+                       float output_norm_eps = 0.f, int block_stride_m = 0,
+                       int block_stride_r = 0) {
+  constexpr size_t smem_size =
+      ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
+       (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
+       sizeof(FwdSmemPlan<NC>) + 15) &
+      ~size_t(15);
+  auto kernel =
+      &attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
+                                     HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
+  static bool attrs_set = false;
+  if (!attrs_set) {
+    if (smem_size > 48 * 1024) {
+      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                           smem_size);
+    }
+    attrs_set = true;
+  }
+  int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
+  cudaLaunchConfig_t config{};
+  config.gridDim = grid;
+  config.blockDim = BLK;
+  config.dynamicSmemBytes = smem_size;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = 1;
+  cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
+                     res_weight, rms_weight, output, N, T, B, block_stride_m,
+                     block_stride_r, rms_eps, output_norm_weight,
+                     output_norm_eps);
+}
+
+}  // namespace fwd_prod_v2
+}  // namespace sm100
+
+void kimi_k3_attn_res(torch::stable::Tensor& prefix,
+                      torch::stable::Tensor const& delta,
+                      torch::stable::Tensor const& blocks,
+                      torch::stable::Tensor const& norm_weight,
+                      torch::stable::Tensor const& qk_weight,
+                      torch::stable::Tensor const& output_norm_weight,
+                      torch::stable::Tensor& output, int64_t num_blocks,
+                      double eps, double output_norm_eps) {
+  int const num_tokens = static_cast<int>(prefix.size(0));
+  int const device = prefix.get_device_index();
+  torch::stable::accelerator::DeviceGuard const device_guard(device);
+  cudaDeviceProp const* properties = get_device_prop();
+  STD_TORCH_CHECK(properties->major == 10,
+                  "Kimi K3 AttnRes requires the SM100 family");
+
+  using namespace sm100::fwd_prod_v2;
+  // Two-source chunks and two resident CTAs are beneficial once setup is
+  // amortized by the long, full eight-block prefill workload.
+  if (num_blocks == 8 && num_tokens >= 4096) {
+    launch_fwd<7168, 2, true, true, true, true>(
+        static_cast<bf16_t const*>(blocks.data_ptr()),
+        static_cast<bf16_t*>(prefix.data_ptr()),
+        static_cast<bf16_t const*>(delta.data_ptr()),
+        static_cast<bf16_t const*>(qk_weight.data_ptr()),
+        static_cast<bf16_t const*>(norm_weight.data_ptr()),
+        static_cast<bf16_t*>(output.data_ptr()),
+        static_cast<int>(num_blocks) + 1, num_tokens, 1,
+        static_cast<float>(eps), properties->multiProcessorCount,
+        get_current_cuda_stream(device),
+        static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
+        static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
+        static_cast<int>(blocks.stride(1)));
+  } else {
+    launch_fwd<7168, 4, false, true, true, true>(
+        static_cast<bf16_t const*>(blocks.data_ptr()),
+        static_cast<bf16_t*>(prefix.data_ptr()),
+        static_cast<bf16_t const*>(delta.data_ptr()),
+        static_cast<bf16_t const*>(qk_weight.data_ptr()),
+        static_cast<bf16_t const*>(norm_weight.data_ptr()),
+        static_cast<bf16_t*>(output.data_ptr()),
+        static_cast<int>(num_blocks) + 1, num_tokens, 1,
+        static_cast<float>(eps), properties->multiProcessorCount,
+        get_current_cuda_stream(device),
+        static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
+        static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
+        static_cast<int>(blocks.stride(1)));
+  }
+  cudaError_t const error = cudaGetLastError();
+  STD_TORCH_CHECK(
+      error == cudaSuccess,
+      "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error));
+}

csrc/libtorch_stable/layernorm_kernels.cu

@@ -249,7 +249,9 @@ void rms_norm(torch::stable::Tensor& out,    // [..., hidden_size]
   int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
 
   // For large num_tokens, use smaller blocks to increase SM concurrency.
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 grid(num_tokens);
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -325,8 +327,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
   /* This kernel is memory-latency bound in many scenarios.
      When num_tokens is large, a smaller block size allows
      for increased block occupancy on CUs and better latency
-     hiding on global mem ops. */
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+     hiding on global mem ops. In batch-invariant mode the block size must
+     not depend on num_tokens, otherwise the same token would use a different
+     reduction width (and thus a different floating-point summation order)
+     across batches; lock it to 1024 to keep results bit-exact. */
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -337,7 +344,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
   auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
   bool offsets_are_multiple_of_vector_width =
       hidden_size % vector_width == 0 && input_stride % vector_width == 0;
-  bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
   const bool has_weight = weight.has_value();
   if (has_weight) {
     auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant(
   int num_tokens = input.numel() / hidden_size;
 
   // For large num_tokens, use smaller blocks to increase SM concurrency.
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 grid(num_tokens);
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(
      When num_tokens is large, a smaller block size allows
      for increased block occupancy on CUs and better latency
      hiding on global mem ops. */
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant(
   auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
   bool ptrs_are_aligned =
       inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
-  bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
   if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
       !batch_invariant_launch) {
     LAUNCH_FUSED_ADD_RMS_NORM(8);

csrc/libtorch_stable/moe/moe_ops.h

@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
                   torch::stable::Tensor& topk_indices,
                   torch::stable::Tensor& token_expert_indices,
                   torch::stable::Tensor& gating_output, bool renormalize,
-                  std::optional<torch::stable::Tensor> bias);
+                  std::optional<torch::stable::Tensor> bias,
+                  std::optional<torch::stable::Tensor> is_padding);
 
 void topk_sigmoid(torch::stable::Tensor& topk_weights,
                   torch::stable::Tensor& topk_indices,
                   torch::stable::Tensor& token_expert_indices,
                   torch::stable::Tensor& gating_output, bool renormalize,
                   std::optional<torch::stable::Tensor> bias,
-                  double routed_scaling_factor);
+                  double routed_scaling_factor,
+                  std::optional<torch::stable::Tensor> is_padding);
 
 void topk_softplus_sqrt(
     torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
     double routed_scaling_factor,
     const std::optional<torch::stable::Tensor>& correction_bias,
     const std::optional<torch::stable::Tensor>& input_ids,
-    const std::optional<torch::stable::Tensor>& tid2eid);
+    const std::optional<torch::stable::Tensor>& tid2eid,
+    const std::optional<torch::stable::Tensor>& is_padding);
 
 void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
              std::optional<torch::stable::Tensor> topk_ids,

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
     const int end_expert,
     const bool renormalize,
     const float* bias,
-    const double routed_scaling_factor)
+    const double routed_scaling_factor,
+    const bool* is_padding)
 {
 
     using cub_kvp = cub::KeyValuePair<int, float>;
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
             const int expert = result_kvp.key;
             const bool node_uses_expert = expert >= start_expert && expert < end_expert;
             const bool should_process_row = row_is_active && node_uses_expert;
+            const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
 
             const int idx = k * block_row + k_idx;
             // Return the unbiased scores for output weights
             output[idx] = inputs_after_softmax[thread_read_offset + expert];
-            indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
-            assert(indices[idx] >= 0);
+            indices[idx] = is_pad_row ? static_cast<IndType>(-1)
+                                       : (should_process_row ? (expert - start_expert) : num_experts);
+            assert(is_pad_row || indices[idx] >= 0);
             source_rows[idx] = k_idx * num_rows + block_row;
             if (renormalize) {
                 selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
     void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
         int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
-        const float* bias, const double routed_scaling_factor)
+        const float* bias, const double routed_scaling_factor, const bool* is_padding)
 {
     static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
                       std::is_same_v<InputType, __half>,
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
             // Add a guard to ignore experts not included by this node
             const bool node_uses_expert = expert >= start_expert && expert < end_expert;
             const bool should_process_row = row_is_active && node_uses_expert;
+            const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
 
             // The lead thread from each sub-group will write out the final results to global memory. (This will be a
             // single) thread per row of the input/output matrices.
             const int idx = k * thread_row + k_idx;
             output[idx] = max_val;
-            indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+            indices[idx] = is_pad_row ? static_cast<IndType>(-1)
+                                       : (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
             source_rows[idx] = k_idx * num_rows + thread_row;
             if (renormalize) {
                 selected_sum += max_val;
@@ -605,7 +610,7 @@ struct TopkConstants
 template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
 void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
     int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
-    const float* bias, const double routed_scaling_factor, cudaStream_t stream)
+    const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
 {
     static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
     using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
 
     dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
     topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
-        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
+        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
 }
 
 #ifndef USE_ROCM
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
                              IndType, InputType, SF>(                         \
         gating_output, nullptr, topk_weights, topk_indices,                   \
         token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,  \
-        bias, routed_scaling_factor, stream);
+        bias, routed_scaling_factor, stream, is_padding);
 #else
   #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                    \
     if (WARP_SIZE == 64) {                                                     \
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
                                IndType, InputType, SF>(                        \
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          bias, routed_scaling_factor, stream);                                                       \
+          bias, routed_scaling_factor, stream, is_padding);                                            \
     } else if (WARP_SIZE == 32) {                                              \
       topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES,       \
                                IndType, InputType, SF>(                        \
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          bias, routed_scaling_factor, stream);                                                       \
+          bias, routed_scaling_factor, stream, is_padding);                                            \
     } else {                                                                   \
       assert(false &&                                                          \
              "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
     const bool renormalize,
     const float* bias,
     const double routed_scaling_factor,
-    cudaStream_t stream) {
+    cudaStream_t stream,
+    const bool* is_padding) {
     static constexpr int WARPS_PER_TB = 4;
     static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
 #ifndef USE_ROCM
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
             }
             moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
                 workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
-                num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
+                num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
         }
     }
 }
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
     int num_tokens, int num_experts, int topk, bool renormalize,
     std::optional<torch::stable::Tensor> bias,
     double routed_scaling_factor,
-    cudaStream_t stream)
+    cudaStream_t stream,
+    std::optional<torch::stable::Tensor> is_padding)
  {
     const float* bias_ptr = nullptr;
     if (bias.has_value()) {
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
       bias_ptr = bias_tensor.const_data_ptr<float>();
     }
 
+    const bool* is_padding_ptr = nullptr;
+    if (is_padding.has_value()) {
+      const torch::stable::Tensor& is_padding_tensor = is_padding.value();
+      STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
+                      "is_padding tensor must be bool");
+      STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
+      STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
+                      "is_padding size mismatch, expected: ", num_tokens);
+      STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
+      is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
+    }
+
     if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
         vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
             reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, routed_scaling_factor, stream);
+            bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
     } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
         vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
             reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, routed_scaling_factor, stream);
+            bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
     } else {
         STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
         vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, routed_scaling_factor, stream);
+            bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
     }
 }
 
@@ -806,7 +825,8 @@ void topk_softmax(
     torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
     torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
     bool renormalize,
-    std::optional<torch::stable::Tensor> bias)
+    std::optional<torch::stable::Tensor> bias,
+    std::optional<torch::stable::Tensor> is_padding)
 {
     const int num_experts = gating_output.size(-1);
     const auto num_tokens = gating_output.numel() / num_experts;
@@ -825,15 +845,15 @@ void topk_softmax(
     if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, 1.0, stream);
+            bias, 1.0, stream, is_padding);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, 1.0, stream);
+            bias, 1.0, stream, is_padding);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, 1.0, stream);
+            bias, 1.0, stream, is_padding);
     } else {
         STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }
@@ -846,7 +866,8 @@ void topk_sigmoid(
     torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
     bool renormalize,
     std::optional<torch::stable::Tensor> bias,
-    double routed_scaling_factor)
+    double routed_scaling_factor,
+    std::optional<torch::stable::Tensor> is_padding)
 {
     const int num_experts = gating_output.size(-1);
     const auto num_tokens = gating_output.numel() / num_experts;
@@ -865,15 +886,15 @@ void topk_sigmoid(
     if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, routed_scaling_factor, stream);
+            bias, routed_scaling_factor, stream, is_padding);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, routed_scaling_factor, stream);
+            bias, routed_scaling_factor, stream, is_padding);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, routed_scaling_factor, stream);
+            bias, routed_scaling_factor, stream, is_padding);
     } else {
         STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

@@ -71,6 +71,80 @@ __device__ __forceinline__ float toFloat(T value) {
   }
 }
 
+#ifndef USE_ROCM
+// Adapted from:
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
+template <typename OutIndType, typename HashIndType>
+__launch_bounds__(128) __global__
+    void dsv4HashTopkSoftplusSqrt(const float* input, float* output,
+                                  OutIndType* indices, int num_rows,
+                                  int num_experts, float routed_scaling_factor,
+                                  const HashIndType* input_ids,
+                                  const HashIndType* tid2eid,
+                                  const bool* is_padding) {
+  const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
+  const int lane = threadIdx.x % 32;
+  if (warp >= num_rows) return;
+  const int64_t token_id = load_index_as_int64(input_ids, warp);
+  const bool is_pad_row = is_padding != nullptr && is_padding[warp];
+
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+  #endif
+  int expert = 0;
+  float weight = 0.f;
+  if (lane < 6 && !is_pad_row) {
+    // only load and calculate for 6 experts
+    expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
+    const float x = input[warp * num_experts + expert];
+    weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
+    if (isnan(weight)) {
+      weight = 0.f;
+    }
+  }
+  float weight_sum = weight;
+  #pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    // sum in warp
+    weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask);
+  }
+
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+  #endif
+  if (lane < 6) {
+    const int offset = warp * 6 + lane;
+    output[offset] =
+        weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
+    indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
+                                  : static_cast<OutIndType>(-1);
+  }
+}
+
+template <typename OutIndType, typename HashIndType>
+void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
+                        int num_rows, int num_experts,
+                        double routed_scaling_factor,
+                        const HashIndType* input_ids,
+                        const HashIndType* tid2eid, cudaStream_t stream,
+                        const bool* is_padding) {
+  if (num_rows == 0) return;
+  auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
+  cudaLaunchConfig_t config = {};
+  config.gridDim = (num_rows + 3) / 4;
+  config.blockDim = 128;
+  config.stream = stream;
+  cudaLaunchAttribute attr;
+  attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attr.val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = &attr;
+  config.numAttrs = 1;
+  const float scale = static_cast<float>(routed_scaling_factor);
+  cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
+                     num_experts, scale, input_ids, tid2eid, is_padding);
+}
+#endif
+
 // ====================== TopK softplus_sqrt things
 // ===============================
 
@@ -99,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
         const int num_rows, IndType* indices, int* source_rows, const int k,
         const int start_expert, const int end_expert, const bool renormalize,
         double routed_scaling_factor, const float* correction_bias,
-        const HashIndType* input_ids, const HashIndType* tid2eid) {
+        const HashIndType* input_ids, const HashIndType* tid2eid,
+        const bool* is_padding) {
   static_assert(std::is_same_v<InputType, float> ||
                     std::is_same_v<InputType, __nv_bfloat16> ||
                     std::is_same_v<InputType, __half>,
@@ -164,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
     return;
   }
   const bool row_is_active = finished ? !finished[thread_row] : true;
+  const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
 
   // We finally start setting up the read pointers for each thread. First, each
   // thread jumps to the start of the row it will read.
@@ -182,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
   cudaGridDependencySynchronize();
 #endif
 
-  // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
-  // to float
-  if constexpr (std::is_same_v<InputType, float>) {
+  if (is_pad_row) {
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii) {
+      row_chunk[ii] = 0.f;
+    }
+  } else if constexpr (std::is_same_v<InputType, float>) {
     using VecType = AlignedArray<float, ELTS_PER_LDG>;
     VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
     const VecType* vec_thread_read_ptr =
@@ -248,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
   if constexpr (USE_HASH) {
     const int64_t token_id = load_index_as_int64(input_ids, thread_row);
     const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
+    if (!is_pad_row) {
 #pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      float val = row_chunk[ii];
-      float val_b = val * beta;
-      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-      row_chunk[ii] = sqrtf(val);
+      for (int ii = 0; ii < VPT; ++ii) {
+        float val = row_chunk[ii];
+        float val_b = val * beta;
+        val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
+        val = sqrtf(val);
+
+        // Dummy/padding tokens can result in NaN values, so
+        // clamp them to 0.0. Note: this clamp could likely be removed if
+        // 'is_padding' is made mandatory
+        if (isnan(val)) {
+          val = 0.f;
+        }
+        row_chunk[ii] = val;
+      }
     }
     float selected_sum = 0.f;
 #pragma unroll
@@ -268,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
                                group_id * THREADS_PER_ROW * ELTS_PER_LDG +
                                local_id;
         if (expert == expert_idx) {
-          indices[idx] = static_cast<IndType>(expert);
+          indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
+                                     : static_cast<IndType>(-1);
           selected_sum += row_chunk[ii];
           break;
         }
@@ -312,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
 #endif
     return;
   } else {
+    if (!is_pad_row) {
 #pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      float val = row_chunk[ii];
-      float val_b = val * beta;
-      // Compute softplus: log(1 + exp(val)) with numerical stability
-      // When val > threshold, softplus(x) ≈ x to avoid exp overflow
-      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-      val = sqrtf(val);
-      if (correction_bias) {
-        const int group_id = ii / ELTS_PER_LDG;
-        const int local_id = ii % ELTS_PER_LDG;
-        const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                               local_id;
-        val = val + correction_bias[expert_idx];
+      for (int ii = 0; ii < VPT; ++ii) {
+        float val = row_chunk[ii];
+        float val_b = val * beta;
+        // Compute softplus: log(1 + exp(val)) with numerical stability
+        // When val > threshold, softplus(x) ≈ x to avoid exp overflow
+        val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
+        val = sqrtf(val);
+        // Dummy/padding tokens can result in NaN values, so
+        // clamp them to 0.0. Note: this clamp could likely be removed if
+        // 'is_padding' is made mandatory
+        if (isnan(val)) {
+          val = 0.f;
+        }
+        if (correction_bias) {
+          const int group_id = ii / ELTS_PER_LDG;
+          const int local_id = ii % ELTS_PER_LDG;
+          const int expert_idx = first_elt_read_by_thread +
+                                 group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                                 local_id;
+          val = val + correction_bias[expert_idx];
+        }
+        row_chunk[ii] = val;
       }
-      row_chunk[ii] = val;
     }
 
     // Original TopK path: find top-k experts by score
@@ -383,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
         // Add a guard to ignore experts not included by this node
         const bool node_uses_expert =
             expert >= start_expert && expert < end_expert;
-        const bool should_process_row = row_is_active && node_uses_expert;
+        const bool should_process_row =
+            row_is_active && node_uses_expert && !is_pad_row;
 
         // The lead thread from each sub-group will write out the final results
         // to global memory. (This will be a single) thread per row of the
         // input/output matrices.
         const int idx = k * thread_row + k_idx;
-        if (correction_bias != nullptr) {
+        if (correction_bias != nullptr && should_process_row) {
           max_val -= correction_bias[expert];
         }
         output[idx] = max_val;
         indices[idx] =
-            should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+            !is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
         source_rows[idx] = k_idx * num_rows + thread_row;
         if (renormalize) {
           selected_sum += max_val;
@@ -477,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
     const int start_expert, const int end_expert, const bool renormalize,
     double routed_scaling_factor, const float* correction_bias,
     const bool use_hash, const HashIndType* input_ids,
-    const HashIndType* tid2eid, cudaStream_t stream) {
+    const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
   static constexpr int BYTES_PER_LDG =
       MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
   using Constants =
@@ -506,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
     cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
                        indices, source_row, k, start_expert, end_expert,
                        renormalize, routed_scaling_factor, correction_bias,
-                       input_ids, tid2eid);
+                       input_ids, tid2eid, is_padding);
 #else
     kernel<<<num_blocks, block_dim, 0, stream>>>(
         input, finished, output, num_rows, indices, source_row, k, start_expert,
         end_expert, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid);
+        input_ids, tid2eid, is_padding);
 #endif
   })
 }
@@ -525,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
         gating_output, nullptr, topk_weights, topk_indices,                    \
         token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,   \
         routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid,  \
-        stream);
+        stream, is_padding);
 #else
   #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
     if (WARP_SIZE == 64) {                                                     \
@@ -534,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
           routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
+          tid2eid, stream, is_padding);                                        \
     } else if (WARP_SIZE == 32) {                                              \
       topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32,      \
                                            MAX_BYTES>(                         \
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
           routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
+          tid2eid, stream, is_padding);                                        \
     } else {                                                                   \
       assert(false &&                                                          \
              "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
@@ -555,7 +654,18 @@ void topkGatingSoftplusSqrtKernelLauncher(
     const int topk, const bool renormalize, double routed_scaling_factor,
     const float* correction_bias, const bool use_hash,
     const HashIndType* input_ids, const HashIndType* tid2eid,
-    cudaStream_t stream) {
+    cudaStream_t stream, const bool* is_padding) {
+#ifndef USE_ROCM
+  if constexpr (std::is_same_v<InputType, float>) {
+    if (use_hash && topk == 6 && renormalize &&
+        (num_experts == 256 || num_experts == 384)) {
+      launchDsv4HashTopk<IndType, HashIndType>(
+          gating_output, topk_weights, topk_indices, num_tokens, num_experts,
+          routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
+      return;
+    }
+  }
+#endif
   static constexpr int WARPS_PER_TB = 4;
   static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
   // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
@@ -650,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
     int num_experts, int topk, bool renormalize, double routed_scaling_factor,
     const std::optional<torch::stable::Tensor>& correction_bias,
     const std::optional<torch::stable::Tensor>& input_ids,
-    const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
+    const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
+    const std::optional<torch::stable::Tensor>& is_padding) {
   const float* bias_ptr = nullptr;
   if (correction_bias.has_value()) {
     bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -659,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
   auto launch = [&](auto* topk_indices_ptr) {
     using OutIndType =
         typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
+
+    const bool* is_padding_ptr = nullptr;
+    if (is_padding.has_value()) {
+      const torch::stable::Tensor& is_padding_tensor = is_padding.value();
+      STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
+                          torch::headeronly::ScalarType::Bool,
+                      "is_padding tensor must be bool");
+      STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
+                      "is_padding tensor must be 1D");
+      STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
+                      "is_padding size mismatch, expected: ", num_tokens);
+      STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
+                      "is_padding tensor must be contiguous");
+      is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
+    }
+
     if (tid2eid.has_value()) {
       STD_TORCH_CHECK(input_ids.has_value(),
                       "input_ids is required for hash MoE");
@@ -673,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
             topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
             num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
             bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
-            tid2eid.value().const_data_ptr<int64_t>(), stream);
+            tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
       } else {
         STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
                         torch::headeronly::ScalarType::Int);
@@ -683,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
             topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
             num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
             bias_ptr, true, input_ids.value().const_data_ptr<int>(),
-            tid2eid.value().const_data_ptr<int>(), stream);
+            tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
       }
     } else {
       vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -691,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
           topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
           num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
           bias_ptr, false, static_cast<const OutIndType*>(nullptr),
-          static_cast<const OutIndType*>(nullptr), stream);
+          static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
     }
   };
 
@@ -715,7 +842,8 @@ void topk_softplus_sqrt(
     bool renormalize, double routed_scaling_factor,
     const std::optional<torch::stable::Tensor>& correction_bias,
     const std::optional<torch::stable::Tensor>& input_ids,
-    const std::optional<torch::stable::Tensor>& tid2eid) {
+    const std::optional<torch::stable::Tensor>& tid2eid,
+    const std::optional<torch::stable::Tensor>& is_padding) {
   const int num_experts = gating_output.size(-1);
   const auto num_tokens = gating_output.numel() / num_experts;
   const int topk = topk_weights.size(-1);
@@ -728,21 +856,22 @@ void topk_softplus_sqrt(
     dispatch_topk_softplus_sqrt_launch<float>(
         gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
         token_expert_indices, num_tokens, num_experts, topk, renormalize,
-        routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
+        routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
+        is_padding);
   } else if (gating_output.scalar_type() ==
              torch::headeronly::ScalarType::Half) {
     dispatch_topk_softplus_sqrt_launch<__half>(
         reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
         topk_weights, topk_indices, token_expert_indices, num_tokens,
         num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
+        input_ids, tid2eid, stream, is_padding);
   } else if (gating_output.scalar_type() ==
              torch::headeronly::ScalarType::BFloat16) {
     dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
         reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
         topk_weights, topk_indices, token_expert_indices, num_tokens,
         num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
+        input_ids, tid2eid, stream, is_padding);
   } else {
     STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
                     gating_output.scalar_type());

csrc/libtorch_stable/moe/torch_bindings.cpp

@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
   m.def(
       "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
-      "bias) -> ()");
+      "bias, Tensor? is_padding) -> ()");
 
   // Apply topk sigmoid to the gating outputs.
   m.def(
       "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, "
-      "Tensor? bias, float routed_scaling_factor) -> ()");
+      "Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
 
   m.def(
       "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, float "
       "routed_scaling_factor, Tensor? "
-      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
+      "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
 
   // Calculate the result of moe by summing up the partial results
   // from all selected experts. topk_ids/expert_map are optional and, when

csrc/libtorch_stable/ops.h

@@ -315,6 +315,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     std::optional<torch::stable::Tensor> index_q_out,
     const std::string& kv_cache_dtype, bool skip_index_branch);
 
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+void kimi_k3_attn_res(torch::stable::Tensor& prefix,
+                      torch::stable::Tensor const& delta,
+                      torch::stable::Tensor const& blocks,
+                      torch::stable::Tensor const& norm_weight,
+                      torch::stable::Tensor const& qk_weight,
+                      torch::stable::Tensor const& output_norm_weight,
+                      torch::stable::Tensor& output, int64_t num_blocks,
+                      double eps, double output_norm_eps);
+#endif
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -527,9 +538,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
                                                     // 656]
     torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& seq_lens,          // [BATCH]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
 
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(

csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -2,6 +2,7 @@
 #include "../../torch_utils.h"
 
 #include "../../dispatch_utils.h"
+#include "../../../core/batch_invariant.hpp"
 #include "layernorm_utils.cuh"
 #include "quant_conversions.cuh"
 
@@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch(
   auto num_tokens = input.numel() / hidden_size;
 
   dim3 grid(num_tokens);
-  const int max_block_size = (num_tokens <= 256) ? 512 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());

csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu

@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
     // AWQ zeros: (size_k // group_size, size_n // 8)
     const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
     int32_t group_size) {
-  int32_t val =
-      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
-  int32_t zero =
-      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
-             blockIdx.y];
+  // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
+  // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
+  // same row (stride 1) instead of striding across rows (stride size_n/8).
+  int col = blockIdx.y * 32 + threadIdx.x;
+  if (col >= size_n / 8) return;
+  (void)size_k;
+
+  int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
+  int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
   int32_t new_val = 0;
 
 #pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
     zero >>= 4;
   }
 
-  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
+  output[blockIdx.x * (size_n / 8) + col] = new_val;
 }
 
 torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
                     "qweight.size(0) % qzeros.size(0) != 0");
     STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
 
-    dim3 blocks(size_k / 32, size_n / 8);
+    dim3 blocks(size_k, (size_n / 8 + 31) / 32);
     marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
         reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
         reinterpret_cast<int32_t*>(output.mutable_data_ptr()),

csrc/libtorch_stable/torch_bindings.cpp

@@ -468,6 +468,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
       "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
 
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+  ops.def(
+      "kimi_k3_attn_res("
+      "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
+      "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
+      "int num_blocks, float eps, float output_norm_eps) -> ()");
+#endif
+
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -693,6 +701,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
 #endif
   ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
            TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+  ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
+#endif
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",
@@ -847,8 +858,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
 
   ops.def(
       "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
-      "batch_size) -> ()");
+      "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
+      "seq_starts) -> ()");
 
   ops.def(
       "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "

csrc/quickreduce/quick_reduce.h

@@ -22,17 +22,25 @@ template <typename AllReduceKernel, typename T>
 __global__ __quickreduce_launch_bounds_two_shot__ static void
 allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
                             int rank, uint8_t** dbuffer_list,
-                            uint32_t data_offset, uint32_t flag_color,
+                            uint32_t data_offset, uint32_t* d_flag_counters,
                             int64_t data_size_per_phase) {
   int block = blockIdx.x;
   int grid = gridDim.x;
 
+  // Load this block's counter from device memory and advance it on-device,
+  // so the color keeps changing across graph replays instead of being frozen.
+  uint32_t flag_color = d_flag_counters[blockIdx.x];
+
   while (block < num_blocks) {
     AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
                          flag_color, data_size_per_phase);
     block += grid;
     flag_color++;
   }
+  // All threads compute the same final value; one writer per block is enough.
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    d_flag_counters[blockIdx.x] = flag_color;
+  }
 }
 
 #define TWOSHOT_DISPATCH(__codec)                                           \
@@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
     hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
                        dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
                        num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color, this->kMaxProblemSize);                  \
+                       d_flag_counters, this->kMaxProblemSize);             \
   } else if (world_size == 4) {                                             \
     using LineCodec = __codec<T, 4>;                                        \
     using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
     hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
                        dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
                        num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color, this->kMaxProblemSize);                  \
+                       d_flag_counters, this->kMaxProblemSize);             \
   } else if (world_size == 8) {                                             \
     using LineCodec = __codec<T, 8>;                                        \
     using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
     hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
                        dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
                        num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color, this->kMaxProblemSize);                  \
+                       d_flag_counters, this->kMaxProblemSize);             \
   }
 
 // INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
@@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
     hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
                        dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
                        num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color, this->kMaxProblemSize);                  \
+                       d_flag_counters, this->kMaxProblemSize);             \
   } else {                                                                  \
     throw std::runtime_error(                                               \
         "INT3 quick all-reduce is only supported for world_size == 2 "      \
@@ -94,7 +102,7 @@ struct DeviceComms {
   static int constexpr kMaxWorldSize = 8;
 
   bool initialized = false;
-  uint32_t flag_color = 1;
+  uint32_t* d_flag_counters = nullptr;
   int world_size;
   int rank;
 
@@ -128,6 +136,16 @@ struct DeviceComms {
     // Clear the flags buffer.
     HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
 
+    // One flag-color counter per block, advanced by the kernel. Start at 1
+    // to stay clear of the flags buffer we just zeroed.
+    HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
+    {
+      std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
+      HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
+                          kMaxNumBlocks * sizeof(uint32_t),
+                          hipMemcpyHostToDevice));
+    }
+
     // Device-side list of IPC buffers.
     buffer_list.resize(world_size);
     HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
@@ -144,6 +162,12 @@ struct DeviceComms {
   hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
 
   void destroy() {
+    // Allocated before `initialized` flips true, so free it on its own guard
+    // to avoid a leak if init fails partway through.
+    if (d_flag_counters) {
+      HIP_CHECK(hipFree(d_flag_counters));
+      d_flag_counters = nullptr;
+    }
     if (initialized) {
       for (int i = 0; i < world_size; i++) {
         if (i != rank) {
@@ -211,8 +235,6 @@ struct DeviceComms {
         break;
     }
     HIP_CHECK(cudaGetLastError());
-    // Rotate the flag color.
-    flag_color += divceil(N, grid);
   }
 };
 

csrc/spinloop.cpp

@@ -7,7 +7,7 @@ extern "C" {
 
 #if defined(__i386__) || defined(__x86_64__)
   #include <cpuid.h>
-  #include <mwaitxintrin.h>
+  #include <x86intrin.h>
 #endif
 
 #if defined(CLOCK_MONOTONIC_RAW)

docker/Dockerfile

@@ -22,9 +22,13 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================
 
-ARG CUDA_VERSION=13.0.2
+ARG CUDA_VERSION=13.0.3
 ARG PYTHON_VERSION=3.12
 ARG UBUNTU_VERSION=22.04
+# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
+# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
+# the default NCCL version shipped with the base image.
+ARG NCCL_VERSION=2.30.7
 
 # By parameterizing the base images, we allow third-party to use their own
 # base images. One use case is hermetic builds with base images stored in
@@ -477,10 +481,17 @@ WORKDIR /workspace
 # Build DeepEP wheels
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
 # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
-ARG DEEPEP_COMMIT_HASH=73b6ea4
+ARG DEEPEP_COMMIT_HASH=d4f41e4e93
 ARG NVSHMEM_VER
+ARG NCCL_VERSION
 RUN --mount=type=cache,target=/opt/uv/cache \
     mkdir -p /tmp/ep_kernels_workspace/dist && \
+    CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
+    if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
+        echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
+            > /tmp/nccl-override.txt && \
+        export UV_OVERRIDE=/tmp/nccl-override.txt; \
+    fi && \
     export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
     /tmp/install_python_libraries.sh \
         --workspace /tmp/ep_kernels_workspace \
@@ -644,6 +655,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
 
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG NCCL_VERSION
 ARG DEADSNAKES_MIRROR_URL
 ARG DEADSNAKES_GPGKEY_URL
 ARG GET_PIP_URL
@@ -696,7 +708,6 @@ RUN apt-get update -y \
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
-    CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \
     apt-get update -y && \
     apt-get install -y --no-install-recommends --allow-change-held-packages \
         cuda-nvcc-${CUDA_VERSION_DASH} \
@@ -709,12 +720,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         libnuma-dev \
         # numactl CLI for NUMA binding at runtime
         numactl && \
-    # Fixes nccl_allocator requiring nccl.h at runtime
-    # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
-    # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
-    # so we pin the version to match our CUDA version
-    NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \
-    apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
     rm -rf /var/lib/apt/lists/*
 
 # Install uv for faster pip installs
@@ -734,6 +739,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
     && chgrp -R 0 /opt/uv \
     && chmod -R g+rwX,a+rX /opt/uv
 
+# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
+# an older version as a transitive dep; this override forces uv to use our
+# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
+RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
+    if [ "$CUDA_MAJOR" -ge 13 ]; then \
+        echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
+            > /etc/uv-overrides.txt; \
+    else \
+        touch /etc/uv-overrides.txt; \
+    fi
+ENV UV_OVERRIDE=/etc/uv-overrides.txt
+
 # ----------------------------------------------------------------------
 # Non-root support (opt-in)
 # ----------------------------------------------------------------------
@@ -793,7 +810,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.14
+ARG FLASHINFER_VERSION=0.6.15.post1
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/Dockerfile.rocm

@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
-# RIXL/UCX build stages
-FROM base AS build_rixl
-ARG RIXL_BRANCH="39be1de8"
-ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="bfb51733"
+# NIXL/UCX build stages
+FROM base AS build_nixl
+ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
+ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
+ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
 ARG UCX_REPO="https://github.com/openucx/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
-ENV RIXL_HOME=/usr/local/rixl
-ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
+ENV NIXL_HOME=/usr/local/nixl
 
-# RIXL build system dependences and RDMA support
+# NIXL build system dependencies and RDMA support
 RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
     libgrpc-dev \
     libgrpc++-dev \
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
     && rm -rf /var/lib/apt/lists/*
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system meson auditwheel patchelf tomlkit
+    uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
+        auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
 
 RUN --mount=type=cache,target=/root/.cache/ccache \
     cd /usr/local/src && \
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
 ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
 
 RUN --mount=type=cache,target=/root/.cache/ccache \
-    git clone ${RIXL_REPO} /opt/rixl && \
-    cd /opt/rixl && \
-    git checkout ${RIXL_BRANCH} && \
+    git clone ${NIXL_REPO} /opt/nixl && \
+    cd /opt/nixl && \
+    git checkout ${NIXL_BRANCH} && \
     CC="ccache gcc" CXX="ccache g++" \
-    meson setup build --prefix=${RIXL_HOME} \
+    meson setup build --prefix=${NIXL_HOME} \
                      -Ducx_path=${UCX_HOME} \
-                     -Drocm_path=${ROCM_PATH} && \
+                     -Dwheel_variant=rocm \
+                     -Dbuild_tests=false \
+                     -Dbuild_examples=false && \
     cd build && \
     ninja -j$(nproc) && \
-    ninja install
+    ninja install && \
+    echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
+        > /etc/ld.so.conf.d/nixl.conf && \
+    echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
+        >> /etc/ld.so.conf.d/nixl.conf && \
+    ldconfig
 
-# Generate RIXL wheel
+# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
+# so configure the ROCm wheel variant directly through Meson.
 # Exclude libcore and libpull from auditwheel: transitive dependencies
 # that are not shipped in the wheel and vary across base images.
-RUN cd /opt/rixl && \
-    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
-        contrib/build-wheel.sh && \
-    mkdir -p /app/install && \
-    _ucx_install_dir=${UCX_HOME} \
-    ./contrib/build-wheel.sh \
-        --output-dir /app/install \
-        --rocm-dir ${ROCM_PATH} \
+RUN cd /opt/nixl && \
+    ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
+    CC="ccache gcc" CXX="ccache g++" \
+    uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
+        --python ${PYTHON_VERSION} \
+        -Csetup-args=-Ducx_path=${UCX_HOME} \
+        -Csetup-args=-Dwheel_variant=rocm \
+        -Csetup-args=-Dbuild_tests=false \
+        -Csetup-args=-Dbuild_examples=false && \
+    mkdir -p /tmp/nixl_wheels/repaired /app/install && \
+    auditwheel repair \
+        --exclude 'libamdhip64*' \
+        --exclude 'libcore*' \
+        --exclude 'libpull*' \
+        /tmp/nixl_wheels/nixl_rocm*.whl \
+        --plat manylinux_2_34_$(uname -m) \
+        --wheel-dir /tmp/nixl_wheels/repaired && \
+    ./contrib/wheel_add_ucx_plugins.py \
         --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
-        --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
+        --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
+        /tmp/nixl_wheels/repaired/*.whl && \
+    cp /tmp/nixl_wheels/repaired/*.whl /app/install
 
 # ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
 # invalidate the slow ROCShmem build.
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
     ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
 fi
 
-# Install RIXL + DeepEP wheels.
-RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
+# Install NIXL + DeepEP wheels.
+RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
     --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
-    uv pip install --system /rixl_install/*.whl /deep_install/*.whl
+    uv pip install --system /nixl_install/*.whl /deep_install/*.whl
 
 # Copy ROCShmem runtime libraries.
 COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -724,6 +744,8 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+ENV UCX_RMA_PPLN_ENABLE=y
+ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
 
 # ROCm profiler limits workaround.
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -796,9 +818,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
     && pip uninstall -y vllm \
     && uv pip install --system *.whl
 
-# Install RIXL wheel
-RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
-    uv pip install --system /rixl_install/*.whl
+# Install NIXL ROCm wheel
+RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
+    uv pip install --system /nixl_install/*.whl
 
 ARG COMMON_WORKDIR
 ARG BASE_IMAGE
@@ -813,6 +835,8 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
 # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+ENV UCX_RMA_PPLN_ENABLE=y
+ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
 
 ENV TOKENIZERS_PARALLELISM=false
 

docker/Dockerfile.rocm_base

@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.16.post3"
+ARG AITER_BRANCH="v0.1.16.post5"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
 ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
 ENV AITER_ROCM_ARCH=gfx942;gfx950
-ENV MORI_GPU_ARCHS=gfx942;gfx950
+# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
 
 # Required for RCCL in ROCm7.1
 ENV HSA_NO_SCRATCH_RECLAIM=1

docker/Dockerfile.s390x

@@ -86,6 +86,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p /tmp/hf-xet/dist && \
     cp dist/*.whl /tmp/hf-xet/dist/
 
+# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
+# llvmlite v0.47 does not support; only SystemZ target is needed).
+FROM base AS llvm20-build
+ARG LLVM_VERSION=20.1.8
+WORKDIR /tmp
+RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
+    curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
+    tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
+    cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
+        -DCMAKE_BUILD_TYPE=Release \
+        -DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
+        -DLLVM_TARGETS_TO_BUILD="SystemZ" \
+        -DLLVM_ENABLE_RTTI=ON \
+        -DLLVM_BUILD_TOOLS=OFF \
+        -DLLVM_BUILD_UTILS=ON \
+        -DLLVM_BUILD_EXAMPLES=OFF \
+        -DLLVM_BUILD_TESTS=OFF \
+        -DLLVM_INCLUDE_TESTS=OFF \
+        -DLLVM_INCLUDE_EXAMPLES=OFF \
+        -DLLVM_INCLUDE_BENCHMARKS=OFF && \
+    ninja -C build install && \
+    rm -rf build llvm-project-${LLVM_VERSION}.src*
+
 # Build numba
 FROM python-install AS numba-builder
 
@@ -96,11 +119,13 @@ WORKDIR /tmp
 
 # Clone all required dependencies
 RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
     microdnf install ninja-build gcc gcc-c++ -y && \
     git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
     git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
     cd llvmlite && \
     uv pip install 'cmake<4' 'setuptools<70' numpy && \
+    CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
     python setup.py bdist_wheel && \
     cd ../numba && \
     if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -158,7 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
      NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
      OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
      uv pip install -v \
-        $ARROW_WHL_FILE \
         $VISION_WHL_FILE \
         $HF_XET_WHL_FILE \
         $LLVM_WHL_FILE \

docker/ci-rocm.hcl

@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
 }
 
 # Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
-# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
+# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
 # Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
 # stage and is irrelevant when building --target ci_base itself.
 variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
 # Upstream dependency commit pins -- extracted from Dockerfile.rocm by
 # ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
 # functions produce no entries when the variable is empty.
-variable "RIXL_BRANCH" {
+variable "NIXL_BRANCH" {
   default = ""
 }
 
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
   default = ""
 }
 
-variable "RIXL_CACHE_KEY" {
+variable "NIXL_CACHE_KEY" {
   default = ""
 }
 
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
   ])
 }
 
-# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
+# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
 # These stages are pinned to specific upstream commit hashes, so cache keys use
 # those hashes rather than the Buildkite commit. This means the cache persists
 # across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
 function "get_cache_from_rocm_deps" {
   params = []
   result = compact([
-    RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
+    NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
     ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
     DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
   ])
 }
 
-function "get_cache_to_rocm_rixl" {
+function "get_cache_to_rocm_nixl" {
   params = []
   result = compact([
-    RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
+    NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
   ])
 }
 
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
 # in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
 # (e.g., requirements change), these stages are cache hits if their upstream
 # pins haven't changed -- saving ~35min of compilation.
-target "rixl-rocm-ci" {
+target "nixl-rocm-ci" {
   inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_rixl"
+  target     = "build_nixl"
   cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_rixl()
+  cache-to   = get_cache_to_rocm_nixl()
   output     = ["type=cacheonly"]
 }
 
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
   output     = ["type=cacheonly"]
 }
 
-# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
+# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
 # Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
 # files drifts from the remote image label. Per-PR builds then pull the result
 # as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
       CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
       CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
     ]),
-    # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
+    # Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
     # are cache hits even when ci_base itself needs rebuilding.
     get_cache_from_rocm_deps(),
   )
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
 # Group for ci_base builds -- exports dependency stage caches alongside the
 # ci_base image so future rebuilds can reuse them independently.
 group "ci-base-rocm-ci-with-deps" {
-  targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
+  targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
 }