type ignore

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/intel_jobs/engine_intel.yaml

@@ -0,0 +1,21 @@
+group: Engine Intel
+depends_on:
+  - image-build-xpu
+steps:
+- label: Engine (1 GPU)
+  timeout_in_minutes: 30
+  device: intel_gpu
+  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/engine/
+    - tests/v1/engine/
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'

.buildkite/intel_jobs/kernels_intel.yaml

@@ -0,0 +1,21 @@
+group: Kernels Intel
+depends_on: 
+  - image-build-xpu
+steps:
+- label: vLLM IR Tests
+  timeout_in_minutes: 30
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/ir
+    - vllm/kernels
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s kernels/ir'

.buildkite/intel_jobs/lora_intel.yaml

@@ -0,0 +1,130 @@
+group: LoRA Intel
+depends_on:
+  - image-build-xpu
+steps:
+- label: LoRA Runtime + Utils
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s lora/test_layers.py &&
+      pytest -v -s lora/test_lora_checkpoints.py &&
+      pytest -v -s lora/test_lora_functions.py &&
+      pytest -v -s lora/test_lora_huggingface.py &&
+      pytest -v -s lora/test_lora_manager.py &&
+      pytest -v -s lora/test_lora_utils.py &&
+      pytest -v -s lora/test_peft_helper.py &&
+      pytest -v -s lora/test_resolver.py &&
+      pytest -v -s lora/test_utils.py &&
+      (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
+      (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
+
+- label: LoRA Fused/MoE Kernels
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
+      pytest -v -s lora/test_moe_lora_align_sum.py'
+
+- label: LoRA Punica Kernels
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      set -o pipefail &&
+      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
+
+- label: LoRA Punica FP8/XPU Ops
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s lora/test_punica_ops_fp8.py &&
+      pytest -v -s lora/test_punica_xpu_ops.py'
+
+- label: LoRA Models
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
+      pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
+      pytest -v -s lora/test_qwen35_densemodel_lora.py &&
+      pytest -v -s lora/test_transformers_model.py'
+
+- label: LoRA Multimodal
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/lora
+    - tests/lora
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s lora/test_default_mm_loras.py && 
+      pytest -v -s lora/test_qwen3_unembed.py &&
+      pytest -v -s lora/test_qwenvl.py &&
+      pytest -v -s lora/test_whisper.py'

.buildkite/intel_jobs/misc_intel.yaml

@@ -0,0 +1,55 @@
+group: Miscellaneous Intel
+depends_on:
+  - image-build-xpu
+steps:
+- label: V1 Core + KV + Metrics
+  timeout_in_minutes: 30
+  device: intel_gpu
+  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/v1/core
+    - tests/v1/executor
+    - tests/v1/kv_offload
+    - tests/v1/worker
+    - tests/v1/kv_connector/unit
+    - tests/v1/metrics
+    - tests/entrypoints/openai/correctness/test_lmeval.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'pip install -r requirements/kv_connectors.txt &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      cd tests &&
+      pytest -v -s v1/executor'
+
+- label: V1 Sample + Logits
+  timeout_in_minutes: 30
+  device: intel_gpu
+  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/v1/sample
+    - tests/v1/logits_processors
+    - tests/v1/test_oracle.py
+    - tests/v1/test_request.py
+    - tests/v1/test_outputs.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      cd tests &&
+      pytest -v -s v1/logits_processors &&
+      pytest -v -s v1/test_oracle.py &&
+      pytest -v -s v1/test_request.py &&
+      pytest -v -s v1/test_outputs.py'

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

@@ -240,7 +240,6 @@ fi
 cleanup_docker
 
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
-aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
 
 # --- Build or pull test image ---
 IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
@@ -283,6 +282,7 @@ docker run \
     --ipc=host \
     --privileged \
     -v /dev/dri/by-path:/dev/dri/by-path \
+    -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
     --entrypoint="" \
     -e "HF_TOKEN=${HF_TOKEN:-}" \
     -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+set -euxo pipefail
+
+# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
+THRESHOLD=${1:-0.8}
+NUM_Q=${2:-1319}
+PORT=${3:-8050}
+OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
+mkdir -p "${OUT_DIR}"
+
+wait_for_server() {
+  local port=$1
+  timeout 600 bash -c '
+    until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
+      sleep 1
+    done'
+}
+
+MODEL="Qwen/Qwen3-30B-A3B-FP8"
+BACK="allgather_reducescatter"
+
+cleanup() {
+  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
+    kill "${SERVER_PID}" 2>/dev/null || true
+    for _ in {1..20}; do
+      kill -0 "${SERVER_PID}" 2>/dev/null || break
+      sleep 0.5
+    done
+    kill -9 "${SERVER_PID}" 2>/dev/null || true
+  fi
+}
+trap cleanup EXIT
+
+VLLM_DEEP_GEMM_WARMUP=skip \
+vllm serve "$MODEL" \
+--enforce-eager \
+--data-parallel-size 4 \
+--enable-expert-parallel \
+--enable-eplb \
+--all2all-backend "$BACK" \
+--eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \
+--trust-remote-code \
+--max-model-len 2048 \
+--port "$PORT" &
+SERVER_PID=$!
+wait_for_server "$PORT"
+
+TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
+OUT="${OUT_DIR}/${TAG}_${BACK}.json"
+python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
+python3 - <<PY
+import json; acc=json.load(open('${OUT}'))['accuracy']
+print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
+assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}"
+PY

.buildkite/test_areas/e2e_integration.yaml

@@ -29,6 +29,15 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
 
+- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
+  timeout_in_minutes: 60
+  device: h100
+  optional: true
+  num_devices: 4
+  working_dir: "/vllm-workspace"
+  commands:
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
+
 - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
   timeout_in_minutes: 60
   device: h100

.buildkite/test_areas/spec_decode.yaml

@@ -12,6 +12,17 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
 
+- label: Spec Decode Eagle Nightly B200
+  timeout_in_minutes: 30
+  device: b200
+  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 -k "eagle_correctness"
+
 - label: Spec Decode Speculators + MTP
   timeout_in_minutes: 30
   device: h200_18gb
@@ -23,6 +34,18 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
 
+- label: Spec Decode Speculators + MTP Nightly B200
+  timeout_in_minutes: 30
+  device: b200
+  optional: true
+  source_file_dependencies:
+    - vllm/v1/spec_decode/
+    - vllm/v1/worker/gpu/spec_decode/
+    - vllm/transformers_utils/configs/speculators/
+    - tests/v1/e2e/spec_decode/
+  commands:
+    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
+  
 - label: Spec Decode Ngram + Suffix
   timeout_in_minutes: 30
   device: h200_18gb
@@ -43,6 +66,17 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
 
+- label: Spec Decode Draft Model Nightly B200
+  timeout_in_minutes: 30
+  device: b200
+  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 -k "draft_model or no_sync or batch_inference"
+
 - label: DFlash Speculators Correctness
   timeout_in_minutes: 30
   device: h100

benchmarks/kernels/ir/bench_ir_ops.py

@@ -0,0 +1,378 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Generic benchmark harness for vLLM IR ops.
+
+Usage:
+    python benchmarks/kernels/ir/bench_ir_ops.py
+    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm
+    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul
+    python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph
+    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/
+"""
+
+import argparse
+import contextlib
+import csv
+import dataclasses
+import datetime
+import math
+import os
+import subprocess
+import sys
+import tempfile
+
+# Ensure repo root is on sys.path so `benchmarks` is importable as a package.
+_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+if _REPO_ROOT not in sys.path:
+    sys.path.insert(0, _REPO_ROOT)
+
+# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2
+# directly by the dynamic linker, so Python-level sys.stderr redirect won't
+# catch them).
+_saved_fd = os.dup(2)
+try:
+    with open(os.devnull, "w") as _devnull:
+        os.dup2(_devnull.fileno(), 2)
+        import torch
+
+        import vllm.kernels  # noqa: E402, F401
+finally:
+    os.dup2(_saved_fd, 2)
+    os.close(_saved_fd)
+
+from tqdm import tqdm  # noqa: E402
+
+from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS  # noqa: E402  # isort: skip
+from vllm.ir.op import IrOp  # noqa: E402
+from vllm.platforms import current_platform  # noqa: E402
+from vllm.triton_utils import triton  # noqa: E402
+
+
+@dataclasses.dataclass(frozen=True)
+class BenchConfig:
+    use_cuda_graph: bool = True
+    warmup: int = 25
+    rep: int = 100
+
+
+def _pkg_version(name: str) -> str:
+    from importlib.metadata import PackageNotFoundError, version
+
+    with contextlib.suppress(PackageNotFoundError):
+        return version(name)
+    return "not installed"
+
+
+_METADATA_LABELS = {
+    "timestamp": "Timestamp",
+    "git_commit": "Git commit",
+    "vllm": "vLLM",
+    "pytorch": "PyTorch",
+    "cuda_runtime": "CUDA runtime",
+    "triton": "Triton",
+    "cutlass": "CUTLASS",
+    "helion": "Helion",
+    "device": "Device",
+    "bench_mode": "Bench mode",
+    "warmup": "Warmup",
+    "rep": "Repetitions",
+}
+
+
+def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]:
+    from vllm.collect_env import get_env_info
+
+    env = get_env_info()
+
+    git_sha = "unknown"
+    with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError):
+        git_sha = (
+            subprocess.check_output(
+                ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
+            )
+            .decode()
+            .strip()
+        )
+
+    device_name = current_platform.get_device_name()
+
+    warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)"
+    rep_note = " replays" if cfg.use_cuda_graph else " ms"
+
+    return {
+        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "git_commit": git_sha,
+        "vllm": str(env.vllm_version),
+        "pytorch": str(env.torch_version),
+        "cuda_runtime": str(env.cuda_runtime_version),
+        "triton": triton.__version__,
+        "cutlass": _pkg_version("nvidia-cutlass-dsl"),
+        "helion": _pkg_version("helion"),
+        "device": device_name,
+        "bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager",
+        "warmup": f"{cfg.warmup}{warmup_note}",
+        "rep": f"{cfg.rep}{rep_note}",
+    }
+
+
+def print_metadata(metadata: dict[str, str]):
+    print("=" * 60)
+    for key, val in metadata.items():
+        print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}")
+    print("=" * 60)
+
+
+def _clone_args(args: tuple) -> tuple:
+    return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
+
+
+# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as
+# inplace=True will mutate bench_args across iterations. Both CUDA graph
+# and eager modes will accumulate drift from repeated in-place mutation.
+# We need to re-clone inputs per iteration for inplace ops.
+def _bench_one(fn, args, cfg: BenchConfig) -> float:
+    bench_args = _clone_args(args)
+    bench_fn = lambda: fn(*bench_args)
+
+    if cfg.use_cuda_graph:
+        ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5])
+    else:
+        ms = triton.testing.do_bench(
+            bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5]
+        )
+    return ms * 1000
+
+
+# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
+# the benchmark baseline should be the compiled native (what vLLM runs by
+# default) rather than the uncompiled native implementation.
+def collect_timings(
+    op: IrOp, shape_configs: list[dict], cfg: BenchConfig
+) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
+    def fmt(v) -> str:
+        return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v)
+
+    case_names = [
+        "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
+    ]
+    providers = [n for n, impl in op.impls.items() if impl.supported]
+
+    results: dict[str, dict[str, float]] = {c: {} for c in case_names}
+    for provider in providers:
+        impl = op.impls[provider]
+        desc = f"{op.name} / {provider}"
+        for case_name, kwargs in tqdm(
+            zip(case_names, shape_configs),
+            desc=desc,
+            total=len(case_names),
+            unit=" cases",
+        ):
+            args = op.generate_inputs(**kwargs)
+            if impl.supports_args(*args):
+                results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
+            else:
+                results[case_name][provider] = float("nan")
+
+    return case_names, providers, results
+
+
+def analyze_results(
+    op_name: str,
+    case_names: list[str],
+    providers: list[str],
+    results: dict[str, dict[str, float]],
+) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
+    native_col = "native"
+    non_native = [p for p in providers if p != native_col]
+
+    header_cols = ["case"]
+    for p in providers:
+        header_cols.append(f"{p} (us)")
+    for p in non_native:
+        header_cols.append(f"{p} speedup")
+
+    detail_rows: list[dict[str, str]] = []
+    speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native}
+
+    for case_name in case_names:
+        timings = results[case_name]
+        row: dict[str, str] = {"case": case_name}
+
+        for p in providers:
+            val = timings.get(p, float("nan"))
+            row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a"
+
+        native_us = timings.get(native_col, float("nan"))
+        for p in non_native:
+            p_us = timings.get(p, float("nan"))
+            if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0:
+                speedup = native_us / p_us
+                row[f"{p} speedup"] = f"{speedup:.2f}x"
+                speedup_data[p].append((speedup, case_name))
+            else:
+                row[f"{p} speedup"] = "n/a"
+
+        detail_rows.append(row)
+
+    summary_rows: list[dict[str, str]] = []
+    for p in non_native:
+        entries = speedup_data[p]
+        if not entries:
+            continue
+        speedups = [s for s, _ in entries]
+        geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups))
+        best_val, best_case = max(entries)
+        worst_val, worst_case = min(entries)
+        wins = sum(1 for s in speedups if s > 1.0)
+        losses = sum(1 for s in speedups if s < 1.0)
+        total = len(speedups)
+
+        print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
+        print(f"  geomean speedup: {geomean:.2f}x")
+        print(f"  best:            {best_val:.2f}x  ({best_case})")
+        print(f"  worst:           {worst_val:.2f}x  ({worst_case})")
+
+        summary_rows.append(
+            {
+                "op": op_name,
+                "provider": p,
+                "geomean_speedup": f"{geomean:.2f}",
+                "best_speedup": f"{best_val:.2f}",
+                "best_case": best_case,
+                "worst_speedup": f"{worst_val:.2f}",
+                "worst_case": worst_case,
+                "wins": str(wins),
+                "losses": str(losses),
+                "total": str(total),
+            }
+        )
+
+    return detail_rows, summary_rows, header_cols
+
+
+def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]):
+    with open(path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def save_results(
+    save_dir: str,
+    op_name: str,
+    detail_rows: list[dict[str, str]],
+    header_cols: list[str],
+    all_summary_rows: list[dict[str, str]],
+    metadata: dict[str, str],
+):
+    write_csv(
+        os.path.join(save_dir, f"{op_name}_detail.csv"),
+        detail_rows,
+        header_cols,
+    )
+    if all_summary_rows:
+        write_csv(
+            os.path.join(save_dir, "summary.csv"),
+            all_summary_rows,
+            list(all_summary_rows[0].keys()),
+        )
+    write_csv(
+        os.path.join(save_dir, "metadata.csv"),
+        [metadata],
+        list(metadata.keys()),
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
+    parser.add_argument(
+        "--ops",
+        type=str,
+        default=None,
+        help="Comma-separated list of op names to benchmark (substring match)",
+    )
+    parser.add_argument(
+        "--no-cuda-graph",
+        action="store_true",
+        help="Disable CUDA graph; use do_bench with L2 cache flushing instead",
+    )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=25,
+        help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)",
+    )
+    parser.add_argument(
+        "--rep",
+        type=int,
+        default=100,
+        help="Repetition time in ms (do_bench) or number of graph replays "
+        "(do_bench_cudagraph) (default: 100)",
+    )
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default=None,
+        help="Directory to save results (default: auto-created temp dir)",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = BenchConfig(
+        use_cuda_graph=not args.no_cuda_graph,
+        warmup=args.warmup,
+        rep=args.rep,
+    )
+
+    torch.set_default_device(current_platform.device_type)
+
+    metadata = collect_env_metadata(cfg)
+    print_metadata(metadata)
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    save_dir = args.save_path or os.path.join(
+        tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}"
+    )
+    os.makedirs(save_dir, exist_ok=True)
+
+    op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
+    all_summary_rows: list[dict[str, str]] = []
+
+    for op in IrOp.registry.values():
+        if op_filters and not any(f in op.name for f in op_filters):
+            continue
+        if not op.has_input_generator:
+            print(f"Skipping op '{op.name}': no input generator registered")
+            continue
+        if op.name not in SHAPE_CONFIGS:
+            raise RuntimeError(
+                f"No benchmark shape config for op '{op.name}'. "
+                f"Add it to benchmarks/kernels/ir/shapes.py"
+            )
+
+        case_names, providers, results = collect_timings(
+            op, SHAPE_CONFIGS[op.name], cfg
+        )
+        detail_rows, summary_rows, header_cols = analyze_results(
+            op.name, case_names, providers, results
+        )
+        all_summary_rows.extend(summary_rows)
+
+        save_results(
+            save_dir,
+            op.name,
+            detail_rows,
+            header_cols,
+            all_summary_rows,
+            metadata,
+        )
+
+    print(f"\nResults saved to: {save_dir}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/kernels/ir/shapes.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shape configurations for IR op benchmarks.
+"""
+
+import torch
+
+NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384]
+COMMON_HIDDEN_SIZES = [
+    2048,  # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
+    3072,  # Gemma 7B/9B
+    4096,  # Llama 3 8B, Qwen 3 8B, Mistral 7B
+    5120,  # Llama 4 Scout 17B-16E
+    7168,  # DeepSeek V3
+    8192,  # Llama 3 70B
+    16384,  # Llama 3 405B
+]
+
+# Each entry maps an op name to a list of kwarg dicts that will be passed
+# to that op's registered input generator via op.generate_inputs(**kwargs).
+SHAPE_CONFIGS: dict[str, list[dict]] = {
+    "rms_norm": [
+        {"num_tokens": n, "hidden_size": d, "dtype": dtype}
+        for dtype in [torch.float16, torch.bfloat16, torch.float32]
+        for d in COMMON_HIDDEN_SIZES
+        for n in NUM_TOKENS
+    ],
+}

cmake/cpu_extension.cmake

@@ -161,16 +161,49 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
-    if(RVV_BF16_FOUND)
-        message(STATUS "BF16 extension detected")
-        set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
-        add_compile_definitions(RISCV_BF16_SUPPORT)
-    elseif (RVV_FP16_FOUND)
-        message(WARNING "BF16 functionality is not available")
-        set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
+    # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
+    # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
+    if(NOT DEFINED VLLM_RVV_VLEN)
+        # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
+        if(EXISTS /proc/cpuinfo)
+            file(READ /proc/cpuinfo _cpuinfo)
+            set(_best 0)
+            foreach(_n IN ITEMS 128 256 512 1024)
+                if(_cpuinfo MATCHES "zvl${_n}b")
+                    set(_best ${_n})
+                endif()
+            endforeach()
+            if(_best GREATER 0)
+                set(VLLM_RVV_VLEN ${_best})
+            endif()
+        endif()
+        # If auto-detect failed (no /proc/cpuinfo or no zvl<N>b reported)
+        # but the compiler supports RVV, require explicit specification.
+        if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
+            message(FATAL_ERROR
+                "RISC-V RVV is available but VLEN could not be auto-detected. "
+                "Please specify VLEN explicitly:\n"
+                "  -DVLLM_RVV_VLEN=128   (for VLEN=128 hardware)\n"
+                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)\n"
+                "  -DVLLM_RVV_VLEN=0     (force scalar, no RVV)")
+        endif()
+    endif()
+    if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
+        message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
+        if(RVV_BF16_FOUND)
+            message(STATUS "BF16 extension detected")
+            set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
+            add_compile_definitions(RISCV_BF16_SUPPORT)
+        elseif(RVV_FP16_FOUND)
+            message(WARNING "BF16 functionality is not available")
+            set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
+        else()
+            message(STATUS "compile riscv with scalar (no FP16/BF16)")
+            set(MARCH_FLAGS -march=rv64gc)
+        endif()
     else()
         message(STATUS "compile riscv with scalar")
-        list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
+        set(MARCH_FLAGS -march=rv64gc)
     endif()
     list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
 else()

csrc/cpu/cpu_types_riscv.hpp

@@ -1,832 +1,25 @@
 #ifndef CPU_TYPES_RISCV_HPP
 #define CPU_TYPES_RISCV_HPP
 
-#include <algorithm>
-#include <cmath>
-#include <cstring>
-#include <iostream>
-#include <limits>
-#include <riscv_vector.h>
-#include <torch/all.h>
+// RISC-V Vector (RVV) CPU type definitions for vLLM.
+//
+// Supports multiple VLENs via compile-time dispatch. The compiler defines
+// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
+// maps VLEN to the correct LMUL suffixes, and the impl header provides
+// VLEN-independent class implementations.
+//
+// To add support for a new VLEN, add the LMUL mapping in
+// cpu_types_riscv_defs.hpp (the impl header needs no changes).
 
-// ============================================================================
-// Vector Register Type Definitions (VLEN=128 bits)
-// ============================================================================
-
-typedef vfloat16m1_t fixed_vfloat16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vfloat16m2_t fixed_vfloat16m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-
-typedef vfloat32m1_t fixed_vfloat32m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vfloat32m2_t fixed_vfloat32m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef vfloat32m4_t fixed_vfloat32m4_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-typedef vfloat32m8_t fixed_vfloat32m8_t
-    __attribute__((riscv_rvv_vector_bits(1024)));
-
-typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
-typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
-
-typedef vuint16m1_t fixed_vuint16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vuint16m2_t fixed_vuint16m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef vuint16m4_t fixed_vuint16m4_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-
-#ifdef RISCV_BF16_SUPPORT
-typedef vbfloat16m1_t fixed_vbfloat16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vbfloat16m2_t fixed_vbfloat16m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef vbfloat16m4_t fixed_vbfloat16m4_t
-    __attribute__((riscv_rvv_vector_bits(512)));
+#ifndef __riscv_vector
+  #error "cpu_types_riscv.hpp included in a non-RVV translation unit"
 #endif
 
-namespace vec_op {
-
-#ifdef RISCV_BF16_SUPPORT
-  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
-    AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
-#else
-  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+#ifndef __riscv_v_min_vlen
+  #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
 #endif
 
-#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+#include "cpu_types_riscv_defs.hpp"
+#include "cpu_types_riscv_impl.hpp"
 
-#define FORCE_INLINE __attribute__((always_inline)) inline
-
-namespace {
-template <typename T, T... indexes, typename F>
-constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
-  (f(std::integral_constant<T, indexes>{}), ...);
-};
-}  // namespace
-
-template <typename T, T count, typename F,
-          typename = std::enable_if_t<std::is_invocable_v<F, T>>>
-constexpr void unroll_loop(F&& f) {
-  unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
-}
-
-template <typename T>
-struct Vec {
-  constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
-};
-
-struct FP32Vec8;
-struct FP32Vec16;
-
-// ============================================================================
-// FP16 Implementation
-// ============================================================================
-
-struct FP16Vec8 : public Vec<FP16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_vfloat16m1_t reg;
-
-  explicit FP16Vec8(const void* ptr)
-      : reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
-                                  VEC_ELEM_NUM)) {};
-
-  explicit FP16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(_Float16);
-    __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
-                           VEC_ELEM_NUM);
-  }
-};
-
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_vfloat16m2_t reg;
-
-  explicit FP16Vec16(const void* ptr)
-      : reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
-                                  VEC_ELEM_NUM)) {};
-
-  explicit FP16Vec16(const FP32Vec16& vec);
-
-  void save(void* ptr) const {
-    __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(_Float16);
-    __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
-                           VEC_ELEM_NUM);
-  }
-};
-
-// ============================================================================
-// BF16 Implementation
-// ============================================================================
-
-#ifdef RISCV_BF16_SUPPORT
-
-FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
-  return __riscv_vreinterpret_v_bf16m1_u16m1(v);
-}
-FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
-  return __riscv_vreinterpret_v_bf16m2_u16m2(v);
-}
-FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
-  return __riscv_vreinterpret_v_bf16m4_u16m4(v);
-}
-
-struct BF16Vec8 : public Vec<BF16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_vbfloat16m1_t reg;
-
-  explicit BF16Vec8(const void* ptr)
-      : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
-  explicit BF16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    __riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
-                           bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-};
-
-struct BF16Vec16 : public Vec<BF16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_vbfloat16m2_t reg;
-
-  explicit BF16Vec16(const void* ptr)
-      : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
-  explicit BF16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const {
-    __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    __riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
-                           bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-};
-
-struct BF16Vec32 : public Vec<BF16Vec32> {
-  constexpr static int VEC_ELEM_NUM = 32;
-  fixed_vbfloat16m4_t reg;
-
-  explicit BF16Vec32(const void* ptr)
-      : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
-
-  explicit BF16Vec32(const BF16Vec8& v) {
-    fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
-    fixed_vuint16m4_t u16_combined =
-        __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
-    reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
-  };
-
-  void save(void* ptr) const {
-    __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
-                          elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    __riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
-                           bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-};
-
-#else
-// ============================================================================
-// BF16 Fallback Implementation (FP32 Simulation)
-// ============================================================================
-
-struct BF16Vec8 : public Vec<BF16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_vfloat32m2_t reg_fp32;
-  explicit BF16Vec8(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[8];
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
-  }
-  explicit BF16Vec8(const FP32Vec8&);
-  void save(void* ptr) const {
-    float tmp[8];
-    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save(void* ptr, int elem_num) const {
-    float tmp[8];
-    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[8];
-    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-
-struct BF16Vec16 : public Vec<BF16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_vfloat32m4_t reg_fp32;
-  explicit BF16Vec16(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[16];
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
-  }
-  explicit BF16Vec16(const FP32Vec16&);
-  void save(void* ptr) const {
-    float tmp[16];
-    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save(void* ptr, int elem_num) const {
-    float tmp[16];
-    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[16];
-    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-
-struct BF16Vec32 : public Vec<BF16Vec32> {
-  constexpr static int VEC_ELEM_NUM = 32;
-  fixed_vfloat32m8_t reg_fp32;
-
-  explicit BF16Vec32(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[32];
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
-  }
-
-  explicit BF16Vec32(const BF16Vec8& v) {
-    float tmp_small[8];
-    __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
-    float tmp_large[32];
-    for (int i = 0; i < 4; ++i) {
-      std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
-    }
-    reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
-  }
-
-  void save(void* ptr) const {
-    float tmp[32];
-    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-
-  void save(void* ptr, int elem_num) const {
-    float tmp[32];
-    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[32];
-    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-#endif
-
-// ============================================================================
-// FP32 Implementation
-// ============================================================================
-
-struct FP32Vec4 : public Vec<FP32Vec4> {
-  constexpr static int VEC_ELEM_NUM = 4;
-  fixed_vfloat32m1_t reg;
-  explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4(const float* ptr)
-      : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
-  void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
-  void save(float* ptr, int elem_num) const {
-    __riscv_vse32_v_f32m1(ptr, reg, elem_num);
-  }
-};
-
-struct FP32Vec8 : public Vec<FP32Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_vfloat32m2_t reg;
-
-  explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(const float* ptr)
-      : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
-  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
-  explicit FP32Vec8(const FP16Vec8& v)
-      : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(fixed_vfloat16m1_t v)
-      : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
-
-#ifdef RISCV_BF16_SUPPORT
-  explicit FP32Vec8(fixed_vbfloat16m1_t v)
-      : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(const BF16Vec8& v)
-      : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
-#else
-  explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
-#endif
-
-  float reduce_sum() const {
-    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  FP32Vec8 operator*(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator+(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator-(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator/(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 min(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 max(const FP32Vec8& b) const {
-    return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 abs() const {
-    return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
-    return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
-  }
-  FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
-    return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
-  }
-
-  FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
-    fixed_vfloat32m2_t temp =
-        __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
-    return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
-  }
-
-  void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
-  void save(float* ptr, int elem_num) const {
-    __riscv_vse32_v_f32m2(ptr, reg, elem_num);
-  }
-  void save_strided(float* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(float);
-    __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
-  }
-
-  FP32Vec8 exp() const {
-    const float inv_ln2 = 1.44269504088896341f;
-    fixed_vfloat32m2_t x_scaled =
-        __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
-    fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
-    fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
-
-    fixed_vfloat32m2_t r =
-        __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
-
-    fixed_vfloat32m2_t poly =
-        __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
-
-    fixed_vint32m2_t biased_exp =
-        __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
-    biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
-    fixed_vint32m2_t exponent_bits =
-        __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
-    fixed_vfloat32m2_t scale =
-        __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
-
-    return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 tanh() const {
-    fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
-        __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
-    fixed_vfloat32m2_t x2 =
-        __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
-    FP32Vec8 exp_val = FP32Vec8(x2).exp();
-    fixed_vfloat32m2_t num =
-        __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
-    fixed_vfloat32m2_t den =
-        __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
-    return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 er() const {
-    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
-                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
-    fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
-
-    fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
-        __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
-    t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
-
-    fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
-                                  a4, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
-                                  a3, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
-                                  a2, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
-                                  a1, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
-
-    fixed_vfloat32m2_t exp_val =
-        FP32Vec8(__riscv_vfneg_v_f32m2(
-                     __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
-                     VEC_ELEM_NUM))
-            .exp()
-            .reg;
-    fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
-        __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-
-    vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
-    return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
-  }
-};
-
-struct FP32Vec16 : public Vec<FP32Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_vfloat32m4_t reg;
-
-  explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(const float* ptr)
-      : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
-  explicit FP32Vec16(const FP32Vec8& data)
-      : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
-  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
-  explicit FP32Vec16(const FP16Vec16& v);
-
-#ifdef RISCV_BF16_SUPPORT
-  explicit FP32Vec16(fixed_vbfloat16m2_t v)
-      : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(const BF16Vec16& v)
-      : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
-#else
-  explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
-#endif
-
-  FP32Vec16 operator+(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator-(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator*(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator/(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  float reduce_sum() const {
-    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  float reduce_max() const {
-    fixed_vfloat32m1_t scalar =
-        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
-    scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  float reduce_min() const {
-    fixed_vfloat32m1_t scalar =
-        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
-    scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  template <int group_size>
-  float reduce_sub_sum(int idx) {
-    static_assert(VEC_ELEM_NUM % group_size == 0);
-    const int start = idx * group_size;
-    vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
-    vbool8_t mask = __riscv_vmand_mm_b8(
-        __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
-        __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
-        VEC_ELEM_NUM);
-    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar =
-        __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  };
-
-  FP32Vec16 max(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 min(const FP32Vec16& b) const {
-    return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 abs() const {
-    return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
-    return FP32Vec16(__riscv_vfmin_vv_f32m4(
-        max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
-        VEC_ELEM_NUM));
-  }
-
-  void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
-  void save(float* ptr, int elem_num) const {
-    __riscv_vse32_v_f32m4(ptr, reg, elem_num);
-  }
-  void save_strided(float* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(float);
-    __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
-  }
-
-  FP32Vec16 exp() const {
-    const float inv_ln2 = 1.44269504088896341f;
-    fixed_vfloat32m4_t x_scaled =
-        __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
-    fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
-    fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
-    fixed_vfloat32m4_t r =
-        __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
-
-    fixed_vfloat32m4_t poly =
-        __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
-                                  0.009618129107628f, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
-                                  0.055504108664821f, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
-                                  0.240226506959101f, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
-                                  0.693147180559945f, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
-                                  1.0f, VEC_ELEM_NUM);
-
-    fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
-        __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
-    fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
-        __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
-
-    return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 tanh() const {
-    fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
-        __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
-    FP32Vec16 exp_val =
-        FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
-    return FP32Vec16(__riscv_vfdiv_vv_f32m4(
-        __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
-        __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 er() const {
-    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
-                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
-    fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
-    fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
-        __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
-                               1.0f, VEC_ELEM_NUM),
-        1.0f, VEC_ELEM_NUM);
-
-    fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
-                                  a4, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
-                                  a3, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
-                                  a2, VEC_ELEM_NUM);
-    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
-                                  a1, VEC_ELEM_NUM);
-    poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
-
-    fixed_vfloat32m4_t exp_val =
-        FP32Vec16(__riscv_vfneg_v_f32m4(
-                      __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
-                      VEC_ELEM_NUM))
-            .exp()
-            .reg;
-    fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
-        __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-
-    vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
-    return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
-  }
-};
-
-// ============================================================================
-// Type Traits & Global Helpers
-// ============================================================================
-
-template <typename T>
-struct VecType {
-  using vec_type = void;
-  using vec_t = void;
-};
-
-template <typename T>
-using vec_t = typename VecType<T>::vec_type;
-
-template <>
-struct VecType<float> {
-  using vec_type = FP32Vec8;
-  using vec_t = FP32Vec8;
-};
-template <>
-struct VecType<c10::Half> {
-  using vec_type = FP16Vec8;
-  using vec_t = FP16Vec8;
-};
-template <>
-struct VecType<c10::BFloat16> {
-  using vec_type = BF16Vec8;
-  using vec_t = BF16Vec8;
-};
-
-template <typename T>
-void storeFP32(float v, T* ptr) {
-  *ptr = v;
-}
-template <>
-inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
-  *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
-}
-
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
-}
-inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
-  reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
-}
-inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
-  reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
-}
-inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc = acc.fma(a, b);
-}
-
-#ifdef RISCV_BF16_SUPPORT
-template <>
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  *ptr = static_cast<__bf16>(v);
-};
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
-    : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
-    : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
-#else
-template <>
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  uint32_t val;
-  std::memcpy(&val, &v, 4);
-  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
-}
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
-#endif
-
-inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
-
-}  // namespace vec_op
-
-#ifndef CPU_KERNEL_GUARD_IN
-  #define CPU_KERNEL_GUARD_IN(NAME)
-#endif
-
-#ifndef CPU_KERNEL_GUARD_OUT
-  #define CPU_KERNEL_GUARD_OUT(NAME)
-#endif
-
-#endif  // CPU_TYPES_RISCV_HPP
+#endif  // CPU_TYPES_RISCV_HPP

csrc/cpu/cpu_types_riscv_defs.hpp

@@ -0,0 +1,98 @@
+#ifndef CPU_TYPES_RISCV_DEFS_HPP
+#define CPU_TYPES_RISCV_DEFS_HPP
+
+// VLEN-to-LMUL mapping for RISC-V Vector extension.
+//
+// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
+//   VLEN=128: LMUL_128=m1,  LMUL_256=m2,  LMUL_512=m4,  LMUL_1024=m8
+//   VLEN=256: LMUL_128=mf2, LMUL_256=m1,  LMUL_512=m2,  LMUL_1024=m4
+
+#include <riscv_vector.h>
+
+#if __riscv_v_min_vlen == 128
+  #define LMUL_128 m1
+  #define LMUL_256 m2
+  #define LMUL_512 m4
+  #define LMUL_1024 m8
+  #define BOOL_256 b16
+  #define BOOL_512 b8
+#elif __riscv_v_min_vlen == 256
+  #define LMUL_128 mf2
+  #define LMUL_256 m1
+  #define LMUL_512 m2
+  #define LMUL_1024 m4
+  #define BOOL_256 b32
+  #define BOOL_512 b16
+#else
+  #error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
+#endif
+
+// Token-paste helpers.
+#define _RVV_P2(a, b) a##b
+#define _RVV_P3(a, b, c) a##b##c
+#define _RVV_P4(a, b, c, d) a##b##c##d
+#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
+#define RVVI(base, lmul) _RVV_P2(base, lmul)
+#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
+#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
+// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
+#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
+#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
+
+// ---- Semantic fixed-vector typedefs (named by element count) ----
+
+// float16
+typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+
+// float32
+typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
+    __attribute__((riscv_rvv_vector_bits(1024)));
+
+// int32
+typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+
+// uint16
+typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+
+// bfloat16
+#ifdef RISCV_BF16_SUPPORT
+typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+#endif
+
+// ---- Reduction accumulator type (always m1 = one register of f32) ----
+// Used for scalar reductions; only element [0] is meaningful.
+typedef vfloat32m1_t rvv_f32_accum_t
+    __attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
+
+// ---- Mask types for f32 elements ----
+#if __riscv_v_min_vlen == 128
+typedef vbool16_t rvv_mask_f32x8_t;
+typedef vbool8_t rvv_mask_f32x16_t;
+#elif __riscv_v_min_vlen == 256
+typedef vbool32_t rvv_mask_f32x8_t;
+typedef vbool16_t rvv_mask_f32x16_t;
+#endif
+
+#endif  // CPU_TYPES_RISCV_DEFS_HPP

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -0,0 +1,891 @@
+#ifndef CPU_TYPES_RISCV_IMPL_HPP
+#define CPU_TYPES_RISCV_IMPL_HPP
+
+// Shared implementation of RVV vector-type wrapper classes.
+// This file is VLEN-independent: it uses the semantic type names and
+// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp.
+//
+// DO NOT include this file directly; include cpu_types_riscv.hpp instead.
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <iostream>
+#include <limits>
+#include <torch/all.h>
+namespace vec_op {
+
+#ifdef RISCV_BF16_SUPPORT
+  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+    AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+#else
+  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+#endif
+
+#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+
+#define FORCE_INLINE __attribute__((always_inline)) inline
+
+namespace {
+template <typename T, T... indexes, typename F>
+constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
+  (f(std::integral_constant<T, indexes>{}), ...);
+};
+}  // namespace
+
+template <typename T, T count, typename F,
+          typename = std::enable_if_t<std::is_invocable_v<F, T>>>
+constexpr void unroll_loop(F&& f) {
+  unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
+}
+
+template <typename T>
+struct Vec {
+  constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
+};
+
+struct FP32Vec8;
+struct FP32Vec16;
+
+// ============================================================================
+// FP16 Implementation
+// ============================================================================
+
+struct FP16Vec8 : public Vec<FP16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_fp16x8_t reg;
+
+  explicit FP16Vec8(const void* ptr)
+      : reg(RVVI(__riscv_vle16_v_f16, LMUL_128)(
+            static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
+
+  explicit FP16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
+                                        VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
+                                        elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(_Float16);
+    RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr),
+                                         byte_stride, reg, VEC_ELEM_NUM);
+  }
+};
+
+struct FP16Vec16 : public Vec<FP16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_fp16x16_t reg;
+
+  explicit FP16Vec16(const void* ptr)
+      : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
+            static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
+
+  explicit FP16Vec16(const FP32Vec16& vec);
+
+  void save(void* ptr) const {
+    RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
+                                        VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
+                                        elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(_Float16);
+    RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr),
+                                         byte_stride, reg, VEC_ELEM_NUM);
+  }
+};
+
+// ============================================================================
+// BF16 Implementation
+// ============================================================================
+
+#ifdef RISCV_BF16_SUPPORT
+
+FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
+  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
+}
+FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) {
+  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v);
+}
+FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) {
+  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v);
+}
+
+struct BF16Vec8 : public Vec<BF16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_bf16x8_t reg;
+
+  explicit BF16Vec8(const void* ptr)
+      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16,
+                  LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {};
+  explicit BF16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
+                                         byte_stride, bf16_to_u16(reg),
+                                         VEC_ELEM_NUM);
+  }
+};
+
+struct BF16Vec16 : public Vec<BF16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_bf16x16_t reg;
+
+  explicit BF16Vec16(const void* ptr)
+      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16,
+                  LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
+  explicit BF16Vec16(const FP32Vec16&);
+
+  void save(void* ptr) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
+                                         byte_stride, bf16_to_u16(reg),
+                                         VEC_ELEM_NUM);
+  }
+};
+
+struct BF16Vec32 : public Vec<BF16Vec32> {
+  constexpr static int VEC_ELEM_NUM = 32;
+  fixed_bf16x32_t reg;
+
+  explicit BF16Vec32(const void* ptr)
+      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
+                  LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
+
+  explicit BF16Vec32(const BF16Vec8& v) {
+    fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
+    fixed_u16x32_t u16_combined =
+        RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
+            u16_val, u16_val, u16_val, u16_val);
+    reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
+                LMUL_512)(u16_combined);
+  };
+
+  void save(void* ptr) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
+                                        bf16_to_u16(reg), elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
+                                         byte_stride, bf16_to_u16(reg),
+                                         VEC_ELEM_NUM);
+  }
+};
+
+#else
+// ============================================================================
+// BF16 Fallback Implementation (FP32 Simulation)
+// ============================================================================
+
+struct BF16Vec8 : public Vec<BF16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_fp32x8_t reg_fp32;
+  explicit BF16Vec8(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[8];
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
+  }
+  explicit BF16Vec8(const FP32Vec8&);
+  void save(void* ptr) const {
+    float tmp[8];
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save(void* ptr, int elem_num) const {
+    float tmp[8];
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[8];
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+
+struct BF16Vec16 : public Vec<BF16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_fp32x16_t reg_fp32;
+  explicit BF16Vec16(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[16];
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
+  }
+  explicit BF16Vec16(const FP32Vec16&);
+  void save(void* ptr) const {
+    float tmp[16];
+    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save(void* ptr, int elem_num) const {
+    float tmp[16];
+    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[16];
+    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+
+struct BF16Vec32 : public Vec<BF16Vec32> {
+  constexpr static int VEC_ELEM_NUM = 32;
+  fixed_fp32x32_t reg_fp32;
+
+  explicit BF16Vec32(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[32];
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
+  }
+
+  explicit BF16Vec32(const BF16Vec8& v) {
+    float tmp_small[8];
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
+    float tmp_large[32];
+    for (int i = 0; i < 4; ++i) {
+      std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
+    }
+    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32);
+  }
+
+  void save(void* ptr) const {
+    float tmp[32];
+    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+
+  void save(void* ptr, int elem_num) const {
+    float tmp[32];
+    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[32];
+    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+#endif
+
+// ============================================================================
+// FP32 Implementation
+// ============================================================================
+
+struct FP32Vec4 : public Vec<FP32Vec4> {
+  constexpr static int VEC_ELEM_NUM = 4;
+  fixed_fp32x4_t reg;
+  explicit FP32Vec4(float v)
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4()
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4(const float* ptr)
+      : reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {};
+  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
+  void save(float* ptr) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
+  }
+  void save(float* ptr, int elem_num) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num);
+  }
+};
+
+struct FP32Vec8 : public Vec<FP32Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_fp32x8_t reg;
+
+  explicit FP32Vec8(float v)
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8()
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(const float* ptr)
+      : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {};
+  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
+  explicit FP32Vec8(const FP16Vec8& v)
+      : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(fixed_fp16x8_t v)
+      : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
+
+#ifdef RISCV_BF16_SUPPORT
+  explicit FP32Vec8(fixed_bf16x8_t v)
+      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(const BF16Vec8& v)
+      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
+        };
+#else
+  explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
+#endif
+
+  float reduce_sum() const {
+    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar,
+                                                               VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  FP32Vec8 operator*(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator+(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator-(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator/(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 min(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 max(const FP32Vec8& b) const {
+    return FP32Vec8(
+        RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 abs() const {
+    return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
+    return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num));
+  }
+  FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
+    return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num));
+  }
+
+  FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
+    fixed_fp32x8_t temp =
+        RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM);
+    return FP32Vec8(
+        RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM));
+  }
+
+  void save(float* ptr) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM);
+  }
+  void save(float* ptr, int elem_num) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num);
+  }
+  void save_strided(float* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(float);
+    RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM);
+  }
+
+  FP32Vec8 exp() const {
+    const float inv_ln2 = 1.44269504088896341f;
+    fixed_fp32x8_t x_scaled =
+        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM);
+    fixed_i32x8_t n_int =
+        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
+    fixed_fp32x8_t n_float =
+        RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM);
+
+    fixed_fp32x8_t r =
+        RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM);
+
+    fixed_fp32x8_t poly =
+        RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f,
+                                                VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f,
+                                                VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f,
+                                                VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f,
+                                                VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM);
+
+    fixed_i32x8_t biased_exp =
+        RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM);
+    biased_exp =
+        RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM);
+    fixed_i32x8_t exponent_bits =
+        RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM);
+    fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32,
+                                 LMUL_256)(exponent_bits);
+
+    return FP32Vec8(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 tanh() const {
+    fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
+        VEC_ELEM_NUM);
+    fixed_fp32x8_t x2 =
+        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM);
+    FP32Vec8 exp_val = FP32Vec8(x2).exp();
+    fixed_fp32x8_t num =
+        RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
+    fixed_fp32x8_t den =
+        RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
+    return FP32Vec8(
+        RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 er() const {
+    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
+                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
+    fixed_fp32x8_t abs_x =
+        RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM);
+
+    fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+    t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM);
+
+    fixed_fp32x8_t poly =
+        RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM);
+
+    fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)(
+                                          RVVI(__riscv_vfmul_vv_f32, LMUL_256)(
+                                              abs_x, abs_x, VEC_ELEM_NUM),
+                                          VEC_ELEM_NUM))
+                                 .exp()
+                                 .reg;
+    fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+
+    rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)(
+        reg, 0.0f, VEC_ELEM_NUM);
+    return FP32Vec8(
+        RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM));
+  }
+};
+
+struct FP32Vec16 : public Vec<FP32Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_fp32x16_t reg;
+
+  explicit FP32Vec16(float v)
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16()
+      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(const float* ptr)
+      : reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {};
+  explicit FP32Vec16(const FP32Vec8& data)
+      : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
+            data.reg, data.reg)) {};
+  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
+  explicit FP32Vec16(const FP16Vec16& v);
+
+#ifdef RISCV_BF16_SUPPORT
+  explicit FP32Vec16(fixed_bf16x16_t v)
+      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(const BF16Vec16& v)
+      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) {
+        };
+#else
+  explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
+#endif
+
+  FP32Vec16 operator+(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator-(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator*(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator/(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  float reduce_sum() const {
+    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar,
+                                                               VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  float reduce_max() const {
+    rvv_f32_accum_t scalar =
+        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
+    scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar,
+                                                              VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  float reduce_min() const {
+    rvv_f32_accum_t scalar =
+        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
+    scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar,
+                                                              VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  template <int group_size>
+  float reduce_sub_sum(int idx) {
+    static_assert(VEC_ELEM_NUM % group_size == 0);
+    const int start = idx * group_size;
+    auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM);
+    rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)(
+        RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start,
+                                                         VEC_ELEM_NUM),
+        RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)(
+            indices, start + group_size, VEC_ELEM_NUM),
+        VEC_ELEM_NUM);
+    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)(
+        mask, reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  };
+
+  FP32Vec16 max(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 min(const FP32Vec16& b) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 abs() const {
+    return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
+    return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)(
+        max_v.reg,
+        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM),
+        VEC_ELEM_NUM));
+  }
+
+  void save(float* ptr) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM);
+  }
+  void save(float* ptr, int elem_num) const {
+    RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num);
+  }
+  void save_strided(float* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(float);
+    RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM);
+  }
+
+  FP32Vec16 exp() const {
+    const float inv_ln2 = 1.44269504088896341f;
+    fixed_fp32x16_t x_scaled =
+        RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM);
+    fixed_i32x16_t n_int =
+        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
+    fixed_fp32x16_t n_float =
+        RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM);
+    fixed_fp32x16_t r =
+        RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM);
+
+    fixed_fp32x16_t poly =
+        RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
+        0.009618129107628f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
+        0.055504108664821f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
+        0.240226506959101f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
+        0.693147180559945f, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+
+    fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)(
+        RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0,
+        VEC_ELEM_NUM);
+    fixed_fp32x16_t scale =
+        RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)(
+            RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM));
+
+    return FP32Vec16(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 tanh() const {
+    fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
+        VEC_ELEM_NUM);
+    FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)(
+                                      x_clamped, 2.0f, VEC_ELEM_NUM))
+                            .exp();
+    return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(
+        RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
+        RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
+        VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 er() const {
+    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
+                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
+    fixed_fp32x16_t abs_x =
+        RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM);
+    fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+            RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f,
+            VEC_ELEM_NUM),
+        1.0f, VEC_ELEM_NUM);
+
+    fixed_fp32x16_t poly =
+        RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1,
+        VEC_ELEM_NUM);
+    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM);
+
+    fixed_fp32x16_t exp_val =
+        FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(
+                      RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x,
+                                                           VEC_ELEM_NUM),
+                      VEC_ELEM_NUM))
+            .exp()
+            .reg;
+    fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+
+    rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)(
+        reg, 0.0f, VEC_ELEM_NUM);
+    return FP32Vec16(
+        RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM));
+  }
+};
+
+// ============================================================================
+// Type Traits & Global Helpers
+// ============================================================================
+
+template <typename T>
+struct VecType {
+  using vec_type = void;
+  using vec_t = void;
+};
+
+template <typename T>
+using vec_t = typename VecType<T>::vec_type;
+
+template <>
+struct VecType<float> {
+  using vec_type = FP32Vec8;
+  using vec_t = FP32Vec8;
+};
+template <>
+struct VecType<c10::Half> {
+  using vec_type = FP16Vec8;
+  using vec_t = FP16Vec8;
+};
+template <>
+struct VecType<c10::BFloat16> {
+  using vec_type = BF16Vec8;
+  using vec_t = BF16Vec8;
+};
+
+template <typename T>
+void storeFP32(float v, T* ptr) {
+  *ptr = v;
+}
+template <>
+inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
+  *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
+}
+
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
+  reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM);
+}
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
+  reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM);
+}
+inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
+  reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM);
+}
+inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
+  acc = acc.fma(a, b);
+}
+
+#ifdef RISCV_BF16_SUPPORT
+template <>
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+  *ptr = static_cast<__bf16>(v);
+};
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
+    : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) {
+      };
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
+    : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
+      };
+#else
+template <>
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+  uint32_t val;
+  std::memcpy(&val, &v, 4);
+  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
+}
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
+#endif
+
+inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
+
+}  // namespace vec_op
+
+#ifndef CPU_KERNEL_GUARD_IN
+  #define CPU_KERNEL_GUARD_IN(NAME)
+#endif
+
+#ifndef CPU_KERNEL_GUARD_OUT
+  #define CPU_KERNEL_GUARD_OUT(NAME)
+#endif
+
+#endif  // CPU_TYPES_RISCV_IMPL_HPP

csrc/moe/topk_softmax_kernels.cu

@@ -126,7 +126,9 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        const float softmax_val = expf(val - float_max) * normalizing_factor;
+        float softmax_val = expf(val - float_max) * normalizing_factor;
+        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
+        if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
         output[idx] = softmax_val;
     }
 }
@@ -147,7 +149,9 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
+        float sigmoid_val = 1.0f / (1.0f + __expf(-val));
+        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
+        if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
         output[idx] = sigmoid_val;
     }
 }
@@ -442,6 +446,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
       }
     }
 
+    // Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
+    // NaN gating (from degenerate hidden states in CUDA graph padding) causes
+    // softmax to produce all-NaN, which makes the argmax loop always pick
+    // expert 0 for every top-k slot, producing duplicate expert IDs that
+    // crash FlashInfer's three-step MoE sort.
+    // With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii) {
+      if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
+        row_chunk[ii] = 0.f;
+      }
+    }
+
     static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
 
     // If bias is not null, use biased value for selection

docker/Dockerfile

@@ -580,33 +580,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
-# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
-#   SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
-ARG FLASHINFER_VERSION=0.6.7
+ARG FLASHINFER_VERSION=0.6.8.post1
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-    && flashinfer show-config
-
-# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
-# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
-# which fails without internet. This step caches them at build time.
-RUN python3 <<'PYEOF'
-from flashinfer.jit import env as jit_env
-from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
-from flashinfer.artifacts import ArtifactPath, CheckSumHash
-
-download_trtllm_headers(
-    'bmm',
-    jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
-    f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
-    ArtifactPath.TRTLLM_GEN_BMM,
-    get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
-)
-
-print('FlashInfer TRTLLM BMM headers downloaded successfully')
-PYEOF
+    && flashinfer show-config \
+    && flashinfer download-cubin
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES

docker/Dockerfile.nightly_torch

@@ -217,16 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.7
-# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
-# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
-#   SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
+# release version: v0.6.8.post1
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/Dockerfile.rocm

@@ -2,6 +2,11 @@
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
+# AMD NIC backend
+ARG NIC_BACKEND=none
+# AMD AINIC apt repo settings
+ARG AINIC_VERSION=1.117.5
+ARG UBUNTU_CODENAME=jammy
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
@@ -216,6 +221,47 @@ RUN git clone ${DEEPEP_REPO} \
  && git checkout ${DEEPEP_BRANCH} \
  && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
 
+# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
+# not force users to rebuild the long-lived Dockerfile.rocm_base image.
+FROM base AS mori_base
+ARG NIC_BACKEND
+ARG AINIC_VERSION
+ARG UBUNTU_CODENAME
+RUN /bin/bash -lc 'set -euo pipefail; \
+ echo "[MORI] Install MoRI proxy deps"; \
+ pip install --quiet --ignore-installed blinker && \
+ pip install --quiet quart msgpack aiohttp pyzmq; \
+ echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
+  \
+  # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
+  # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
+  case "${NIC_BACKEND}" in \
+    # default: mlx5
+    none) \
+      ;; \
+    # AMD NIC
+    ainic) \
+      apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
+      rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
+      curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
+      echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
+        > /etc/apt/sources.list.d/amdainic.list; \
+      apt-get update && apt-get install -y --no-install-recommends \
+          libionic-dev \
+          ionic-common \
+      ; \
+      rm -rf /var/lib/apt/lists/*; \
+      ;; \
+    # TODO: Add Broadcom bnxt packages/repos here later.
+    # bnxt) \
+    #   echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
+    #   ;; \
+    *) \
+      echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
+      exit 2; \
+      ;; \
+  esac;'
+
 # -----------------------
 # vLLM wheel release build stage (for building distributable wheels)
 # This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -294,7 +340,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # -----------------------
 # Test vLLM image
-FROM base AS test
+FROM mori_base AS test
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -385,7 +431,7 @@ RUN printf '%s\n' \
 
 # -----------------------
 # Final vLLM image
-FROM base AS final
+FROM mori_base AS final
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -422,6 +468,8 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
 
 ARG COMMON_WORKDIR
 ARG BASE_IMAGE
+ARG NIC_BACKEND
+ARG AINIC_VERSION
 
 # Copy over the benchmark scripts as well
 COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -439,7 +487,9 @@ ENV HIP_FORCE_DEV_KERNARG=1
 # Workaround for ROCm profiler limits
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
 ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
+    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
+    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
 
 CMD ["/bin/bash"]
 

docker/Dockerfile.rocm_base

@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="v0.1.10.post3"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
-ARG MORI_BRANCH="2d02c6a9"
+ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
 
 # Sccache configuration (only used in release pipeline)

docker/versions.json

@@ -65,7 +65,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.7"
+      "default": "0.6.8.post1"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/contributing/profiling.md

@@ -206,8 +206,8 @@ Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_cont
 used to profile a section of code.
 
 !!! note
-    The legacy import paths `vllm.utils.cprofile` and `vllm.utils.cprofile_context` are deprecated.
-    Please use `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` instead.
+    The `vllm.utils.profiling` helpers are deprecated and will be removed in
+    `v0.21`. Please use Python's `cProfile` module directly instead.
 
 ### Example usage - decorator
 

docs/design/attention_backends.md

@@ -169,7 +169,7 @@ Priority is **1 = highest** (tried first).
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
@@ -215,7 +215,7 @@ configuration.
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

docs/design/cuda_graphs_multimodal.md

@@ -52,14 +52,14 @@ For each graph replay:
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
 
-### Video inference support (experimental)
+### Video inference support
 
 Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
 
 !!! note
     Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
 
-    Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
+    Mixed inputs (image+video) per prompt are also supported now.
 
 ## Model integration via `SupportsEncoderCudaGraph`
 
@@ -76,6 +76,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
 * `get_input_modality(...)` - return the modality of the inputs.
+* `get_max_frames_per_video()` - return model-specific max frames per video.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -96,7 +97,7 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
 * `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
-* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
 
 ## Usage guide
 
@@ -142,7 +143,6 @@ Enable encoder CUDA Graphs via `compilation_config`:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --limit-mm-per-prompt '{"image": 0}' \
   --compilation-config '{"cudagraph_mm_encoder": true}'
 ```
 
@@ -150,7 +150,6 @@ With explicit budgets:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --limit-mm-per-prompt '{"image": 0}' \
   --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
 ```
 
@@ -169,7 +168,6 @@ compilation_config = {
 
 model = vllm.LLM(
     model="Qwen/Qwen3-VL-32B",
-    limit_mm_per_prompt='{"image": 0}',
     compilation_config=compilation_config,
 )
 ```

docs/design/fusions.md

@@ -31,6 +31,7 @@ or just on the low or high end.
 | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant)                      | `fuse_norm_quant`            | RMSNorm (+residual add) → FP8/FP4 quant        | O1 (conditional)               | 1-4%               | No        | Always       |
 | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant)                      | `fuse_act_quant`             | SiLU+Mul activation → FP8/FP4 quant            | O1 (conditional)               | 1-4%               | No        | Always       |
 | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding)                        | `fuse_act_padding`           | Residual add + RMSNorm → padding               | O1 (ROCm/AITER only)           | TBD                | No        | Always       |
+| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm)                   | `fuse_mla_dual_rms_norm`     | Paired Q + KV RMSNorm → single kernel          | O1 (ROCm/AITER only)           | ~2%                | No        | Always       |
 
 ## Support Matrix
 
@@ -51,6 +52,7 @@ The table below lists the quantization schemes supported by each fusion on each
 | `fuse_norm_quant`            | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | —             | FP8 static, FP8 per-token, FP8 per-group |
 | `fuse_act_quant`             | FP8 static, NVFP4                        | FP8 static, FP8 per-group (128/64)       | FP8 static, FP8 per-group (128/64)       | —             | FP8 per-group                            |
 | `fuse_act_padding`           | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
+| `fuse_mla_dual_rms_norm`     | —                                        | —                                        | —                                        | —             | BF16                                     |
 
 \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support
 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
@@ -381,6 +383,44 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled.
 
 - Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`)
 
+### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`)
+
+!!! info
+    ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention.
+
+!!! note
+    When the native implementation of `rms_norm` is used (the default on CUDA and
+    ROCm for now), Inductor's built-in fusion already handles merging these norms
+    automatically. This explicit pass targets the case where AITER's custom
+    `rms_norm` op is active, which Inductor cannot fuse on its own.
+
+**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm
+operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER,
+reducing kernel launch overhead from 2 launches to 1 per MLA layer.
+
+```text
+# Unfused:
+q_c, kv_lora = split(projected, [q_dim, kv_dim])
+kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
+q_c  = rms_norm(q_c,  q_weight,  eps)
+kv_c = rms_norm(kv_c, kv_weight, eps)
+
+# Fused:
+q_c, kv_lora = split(projected, [q_dim, kv_dim])
+kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
+q_normed, kv_normed = fused_mla_dual_rms_norm(
+    q_c, q_weight, kv_c, kv_weight, eps1, eps2)
+```
+
+Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
+when AITER is available.
+
+**Code locations.**
+
+- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)
+- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)
+- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)
+
 ## See Also
 
 - [Optimization Levels](optimization_levels.md) — high-level presets that set

docs/design/optimization_levels.md

@@ -56,6 +56,7 @@ Fusions:
 - `-cc.pass_config.fuse_norm_quant=True`*
 - `-cc.pass_config.fuse_act_quant=True`*
 - `-cc.pass_config.fuse_act_padding=True`†
+- `-cc.pass_config.fuse_mla_dual_rms_norm=True`†
 
 \* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.</br>
 † These fusions are ROCm-only and require AITER.

docs/features/batch_invariance.md

@@ -104,7 +104,7 @@ for output in outputs:
 Batch invariance has been tested and verified on the following models:
 
 - **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
-- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
+- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
 - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
 - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
 - **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`

docs/features/quantization/online.md

@@ -17,6 +17,9 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
 
 # Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
 llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
+
+# MXFP8 quantization for weights and activations
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8")
 ```
 
 Or with the CLI:
@@ -24,6 +27,7 @@ Or with the CLI:
 ```bash
 vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
 vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
+vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
 ```
 
 ## Supported Schemes
@@ -32,8 +36,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
 | ------ | ------------- | ------------------ | ----- |
 | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
 | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
-
-Support for additional schemes will be added in future versions of vllm.
+| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback |
 
 ## Advanced Configuration
 

docs/features/reasoning_outputs.md

@@ -283,9 +283,7 @@ curl http://localhost:8000/v1/chat/completions \
     "messages": [
       { "role": "user", "content": "9.11 and 9.8, which is greater?" }
     ],
-    "extra_body": {
-      "thinking_token_budget": 10
-    }
+    "thinking_token_budget": 10
   }'
 ```
 

docs/models/pooling_models/README.md

@@ -286,10 +286,10 @@ Pooling models now support token-wise task.
 
 ### Score task
 
-`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
-classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
+equal to 1 can it be used as a scoring model and have its scoring API enabled.
 
 ### Pooling multitask support
 
-Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
+Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.

docs/models/supported_models.md

@@ -560,6 +560,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
 | `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
+| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + I<sup>E+</sup> | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ |
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
 | `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
 | `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |

docs/serving/expert_parallel_deployment.md

@@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
 | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
 | `policy` | The policy type for expert parallel load balancing | `"default"` |
+| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`,  or `null` (auto) | `null` |
 
 For example:
 

examples/offline_inference/vision_language.py

@@ -394,18 +394,24 @@ def run_eagle2_5(questions: list[str], modality: str) -> ModelRequestData:
 def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
         max_num_seqs=5,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         trust_remote_code=True,
     )
 
+    image_placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"
+    video_placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"
+
     if modality == "image":
-        placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -425,6 +431,7 @@ def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "LGAI-EXAONE/EXAONE-4.5-33B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -434,18 +441,23 @@ def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<vision><|image_pad|></vision>"
+    video_placeholder = "<vision><|video_pad|></vision>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|system|>\nYou are a helpful assistant.<|endofturn|>\n"
-            f"<|user|>\n<vision>{placeholder}</vision>"
+            f"<|user|>\n{placeholder}"
             f"{question}<|endofturn|>\n"
             "<|assistant|>\n"
         )
@@ -566,6 +578,7 @@ def run_glm4v(questions: list[str], modality: str) -> ModelRequestData:
 def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "zai-org/GLM-4.1V-9B-Thinking"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -574,14 +587,19 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
     )
 
+    image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+    video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+
     if modality == "image":
-        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -602,6 +620,7 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
 def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "zai-org/GLM-4.5V"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -610,15 +629,20 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
         tensor_parallel_size=4,
     )
 
+    image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+    video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+
     if modality == "image":
-        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -639,6 +663,7 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
 def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "zai-org/GLM-4.5V-FP8"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -647,15 +672,20 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
         tensor_parallel_size=4,
     )
 
+    image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+    video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+
     if modality == "image":
-        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -676,6 +706,7 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
 def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "zai-org/GLM-OCR"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -684,14 +715,19 @@ def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
     )
 
+    image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+    video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+
     if modality == "image":
-        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -772,11 +808,12 @@ def run_hyperclovax_seed_vision(
     model_name = "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B"
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         trust_remote_code=True,
-        max_model_len=8192 if modality == "image" else 16384,
-        limit_mm_per_prompt={modality: 1},
+        max_model_len=16384 if modality in ("video", "image+video") else 8192,
+        limit_mm_per_prompt=mm_limit,
     )
 
     messages = list()
@@ -828,6 +865,29 @@ def run_hyperclovax_seed_vision(
                     }
                 ]
             )
+        elif modality == "image+video":
+            messages.append(
+                [
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image",
+                                "ocr": "",
+                                "lens_keywords": "",
+                                "lens_local_keywords": "",
+                            },
+                            {
+                                "type": "video",
+                            },
+                            {
+                                "type": "text",
+                                "text": question,
+                            },
+                        ],
+                    }
+                ]
+            )
         else:
             raise ValueError(f"Unsupported modality: {modality}")
 
@@ -876,19 +936,25 @@ def run_idefics3(questions: list[str], modality: str) -> ModelRequestData:
 def run_interns1(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "internlm/Intern-S1-mini"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         trust_remote_code=True,
         max_model_len=8192,
         max_num_seqs=2,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
     )
 
+    image_placeholder = "<IMG_CONTEXT>"
+    video_placeholder = "<video>"
+
     if modality == "image":
-        placeholder = "<IMG_CONTEXT>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<video>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + "\n" + video_placeholder
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -909,20 +975,26 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData:
 def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "internlm/Intern-S1-Pro"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         trust_remote_code=True,
         max_model_len=8192,
         max_num_seqs=2,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         enforce_eager=True,
         tensor_parallel_size=4,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -943,17 +1015,23 @@ def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData:
 def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "OpenGVLab/InternVL3-2B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         trust_remote_code=True,
         max_model_len=8192,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<image>"
+    video_placeholder = "<video>"
+
     if modality == "image":
-        placeholder = "<image>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<video>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + "\n" + video_placeholder
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -1010,21 +1088,27 @@ def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData:
 def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Kwai-Keye/Keye-VL-8B-Preview"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=8192,
         trust_remote_code=True,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1041,21 +1125,27 @@ def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_keye_vl1_5(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Kwai-Keye/Keye-VL-1.5-8B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=8192,
         trust_remote_code=True,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1259,22 +1349,26 @@ def run_llava_next_video(questions: list[str], modality: str) -> ModelRequestDat
 
 # LLaVA-OneVision
 def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData:
-    if modality == "video":
-        prompts = [
-            f"<|im_start|>user <video>\n{question}<|im_end|><|im_start|>assistant\n"
-            for question in questions
-        ]
+    image_placeholder = "<image>"
+    video_placeholder = "<video>"
 
-    elif modality == "image":
-        prompts = [
-            f"<|im_start|>user <image>\n{question}<|im_end|><|im_start|>assistant\n"
-            for question in questions
-        ]
+    if modality == "image":
+        placeholder = image_placeholder
+    elif modality == "video":
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + "\n" + video_placeholder
 
+    prompts = [
+        (f"<|im_start|>user {placeholder}\n{question}<|im_end|><|im_start|>assistant\n")
+        for question in questions
+    ]
+
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model="llava-hf/llava-onevision-qwen2-7b-ov-hf",
         max_model_len=16384,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
     return ModelRequestData(
@@ -1307,7 +1401,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
 
 # MiniCPM-V
 def run_minicpmv_base(questions: list[str], modality: str, model_name):
-    assert modality in ["image", "video"]
+    assert modality in ["image", "video", "image+video"]
     # If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
 
     # 2.0
@@ -1329,12 +1423,13 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
     # o2.6: image, video, audio
     # model_name = "openbmb/MiniCPM-o-2_6"
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
         max_num_seqs=2,
         trust_remote_code=True,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
     # NOTE The stop_token_ids are different for various versions of MiniCPM-V
     # 2.0
@@ -1347,17 +1442,22 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
     stop_tokens = ["<|im_end|>", "<|endoftext|>"]
     stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
 
-    modality_placeholder = {
-        "image": "(<image>./</image>)",
-        "video": "(<video>./</video>)",
-    }
+    image_placeholder = "(<image>./</image>)"
+    video_placeholder = "(<video>./</video>)"
+
+    if modality == "image":
+        placeholder = image_placeholder
+    elif modality == "video":
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + "\n" + video_placeholder
 
     prompts = [
         tokenizer.apply_chat_template(
             [
                 {
                     "role": "user",
-                    "content": f"{modality_placeholder[modality]}\n{question}",
+                    "content": f"{placeholder}\n{question}",
                 }
             ],
             tokenize=False,
@@ -1466,20 +1566,24 @@ def run_molmo(questions: list[str], modality: str) -> ModelRequestData:
 def run_molmo2(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "allenai/Molmo2-8B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         trust_remote_code=True,
         dtype="bfloat16",
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
         max_num_batched_tokens=36864,
     )
 
+    image_placeholder = "<|image|>"
+    video_placeholder = "<|video|>"
+
     if modality == "image":
-        placeholder = "<|image|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video|>"
-    else:
-        raise ValueError(f"Unsupported modality for molmo2: {modality}")
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         f"{placeholder}<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1563,19 +1667,25 @@ def run_nvlm_d(questions: list[str], modality: str) -> ModelRequestData:
 def run_openpangu_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "FreedomIntelligence/openPangu-VL-7B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
         max_num_seqs=4,
         trust_remote_code=True,
         enforce_eager=True,
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "[unused19]"
+    video_placeholder = "[unused32]"
+
     if modality == "image":
-        placeholder = "[unused19]"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "[unused32]"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
@@ -1623,18 +1733,25 @@ def run_ovis(questions: list[str], modality: str) -> ModelRequestData:
 def run_ovis2_5(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "AIDC-AI/Ovis2.5-2B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
         max_num_seqs=2,
         trust_remote_code=True,
         dtype="half",
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
+
+    image_placeholder = "<image>"
+    video_placeholder = "<video>"
+
     if modality == "image":
-        placeholder = "<image>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<video>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + "\n" + video_placeholder
 
     prompts = [
         f"<|im_start|>user\n\n{placeholder}\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1846,6 +1963,7 @@ def run_qwen_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Qwen/Qwen2-VL-7B-Instruct"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -1855,18 +1973,23 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
             "min_pixels": 28 * 28,
             "max_pixels": 1280 * 28 * 28,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1883,6 +2006,7 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Qwen/Qwen2.5-VL-3B-Instruct"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -1892,18 +2016,23 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1920,6 +2049,7 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_qwen2_5_omni(questions: list[str], modality: str):
     model_name = "Qwen/Qwen2.5-Omni-7B"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -1929,13 +2059,18 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_bos|><|IMAGE|><|vision_eos|>"
+    video_placeholder = "<|vision_bos|><|VIDEO|><|vision_eos|>"
+
     if modality == "image":
-        placeholder = "<|IMAGE|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|VIDEO|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     default_system = (
         "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
@@ -1946,7 +2081,7 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
     prompts = [
         (
             f"<|im_start|>system\n{default_system}<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_bos|>{placeholder}<|vision_eos|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1962,6 +2097,7 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
 def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Qwen/Qwen3-VL-4B-Instruct"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -1971,18 +2107,23 @@ def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1999,6 +2140,7 @@ def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
 def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "Qwen/Qwen3-VL-30B-A3B-Instruct"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -2008,18 +2150,23 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2190,6 +2337,7 @@ def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
 def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
     model_name = "omni-research/Tarsier2-Recap-7b"
 
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
     engine_args = EngineArgs(
         model=model_name,
         max_model_len=4096,
@@ -2197,18 +2345,23 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
             "architectures": ["Tarsier2ForConditionalGeneration"],
             "model_type": "tarsier2",
         },
-        limit_mm_per_prompt={modality: 1},
+        limit_mm_per_prompt=mm_limit,
     )
 
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
     if modality == "image":
-        placeholder = "<|image_pad|>"
+        placeholder = image_placeholder
     elif modality == "video":
-        placeholder = "<|video_pad|>"
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
+            f"<|im_start|>user\n{placeholder}"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2357,6 +2510,24 @@ def get_multi_modal_input(args):
             "questions": vision_chunk_questions,
         }
 
+    if args.modality == "image+video":
+        image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+        needs_metadata = args.model_type in MODELS_NEED_VIDEO_METADATA
+        video = VideoAsset(name="baby_reading", num_frames=args.num_frames).np_ndarrays
+        metadata = VideoAsset(name="baby_reading", num_frames=args.num_frames).metadata
+        img_video_questions = [
+            "What is shown in the image? What happens in the video?",
+            "Describe both the image and the video content.",
+        ]
+
+        return {
+            "data": {
+                "image": image,
+                "video": ([(video, metadata)] if needs_metadata else video),
+            },
+            "questions": img_video_questions,
+        }
+
     msg = f"Modality {args.modality} is not supported."
     raise ValueError(msg)
 
@@ -2439,7 +2610,7 @@ def parse_args():
         "--modality",
         type=str,
         default="image",
-        choices=["image", "video", "vision_chunk"],
+        choices=["image", "video", "image+video", "vision_chunk"],
         help="Modality of the input.",
     )
     parser.add_argument(
@@ -2546,23 +2717,42 @@ def main(args):
         else req_data.sampling_params
     )
 
+    def _mm_data(data, modality):
+        if modality == "image+video":
+            return {"image": data["image"], "video": data["video"]}
+        return {modality: data}
+
+    def _mm_uuid(uuid, modality):
+        if modality == "image+video":
+            return {"image": uuid, "video": uuid + "v"}
+        return {modality: uuid}
+
+    def _mm_empty(modality):
+        if modality == "image+video":
+            return {"image": None, "video": None}
+        return {modality: None}
+
     assert args.num_prompts > 0
     if args.num_prompts == 1:
         # Single inference
         uuid = "uuid_0"
         inputs = {
             "prompt": prompts[0],
-            "multi_modal_data": {modality: data},
-            "multi_modal_uuids": {modality: uuid},
+            "multi_modal_data": _mm_data(data, modality),
+            "multi_modal_uuids": _mm_uuid(uuid, modality),
         }
         inputs_with_empty_media = {
             "prompt": prompts[0],
-            "multi_modal_data": {modality: None},
-            "multi_modal_uuids": {modality: uuid},
+            "multi_modal_data": _mm_empty(modality),
+            "multi_modal_uuids": _mm_uuid(uuid, modality),
         }
     else:
         # Batch inference
         if args.image_repeat_prob is not None:
+            if modality == "image+video":
+                raise ValueError(
+                    "--image-repeat-prob is not supported for 'image+video' modality"
+                )
             # Repeat images with specified probability of "image_repeat_prob"
             inputs, inputs_with_empty_media = apply_image_repeat(
                 args.image_repeat_prob,
@@ -2572,7 +2762,7 @@ def main(args):
                 modality,
             )
         else:
-            # Use the same image for all prompts
+            # Use the same image/video for all prompts
             inputs = []
             inputs_with_empty_media = []
             for i in range(args.num_prompts):
@@ -2580,15 +2770,15 @@ def main(args):
                 inputs.append(
                     {
                         "prompt": prompts[i % len(prompts)],
-                        "multi_modal_data": {modality: data},
-                        "multi_modal_uuids": {modality: uuid},
+                        "multi_modal_data": _mm_data(data, modality),
+                        "multi_modal_uuids": _mm_uuid(uuid, modality),
                     }
                 )
                 inputs_with_empty_media.append(
                     {
                         "prompt": prompts[i % len(prompts)],
-                        "multi_modal_data": {modality: None},
-                        "multi_modal_uuids": {modality: uuid},
+                        "multi_modal_data": _mm_empty(modality),
+                        "multi_modal_uuids": _mm_uuid(uuid, modality),
                     }
                 )
 

examples/offline_inference/vision_language_multi_image.py

@@ -310,6 +310,38 @@ def load_gemma3(question: str, image_urls: list[str]) -> ModelRequestData:
     )
 
 
+def load_granite4_vision(question: str, image_urls: list[str]) -> ModelRequestData:
+    model_name = "ibm-granite/granite-vision-4.1-4b"
+    engine_args = EngineArgs(
+        model=model_name,
+        max_model_len=4096,
+        max_num_seqs=16,
+        limit_mm_per_prompt={"image": len(image_urls)},
+    )
+
+    placeholders = [{"type": "image", "image": url} for url in image_urls]
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                *placeholders,
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+
+    processor = AutoProcessor.from_pretrained(model_name)
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompt=prompt,
+        image_data=[fetch_image(url) for url in image_urls],
+    )
+
+
 def load_h2ovl(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "h2oai/h2ovl-mississippi-800m"
 
@@ -1487,6 +1519,7 @@ model_example_map = {
     "deepseek_ocr": load_deepseek_ocr,
     "exaone4_5": load_exaone4_5,
     "gemma3": load_gemma3,
+    "granite4_vision": load_granite4_vision,
     "h2ovl_chat": load_h2ovl,
     "hunyuan_vl": load_hunyuan_vl,
     "hyperclovax_seed_vision": load_hyperclovax_seed_vision,

examples/pooling/token_embed/jina_embeddings_v4_offline.py

@@ -4,68 +4,74 @@
 import torch
 
 from vllm import LLM
+from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image
 
-# Initialize model
-model = LLM(
-    model="jinaai/jina-embeddings-v4-vllm-text-matching",
-    runner="pooling",
-    max_model_len=1024,
-    gpu_memory_utilization=0.8,
-)
 
-# Create text prompts
-text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+def main():
+    # Initialize model
+    model = LLM(
+        model="jinaai/jina-embeddings-v4-vllm-text-matching",
+        pooler_config=PoolerConfig(task="token_embed"),
+        runner="pooling",
+        max_model_len=1024,
+        gpu_memory_utilization=0.8,
+    )
 
-text2 = "浜辺に沈む美しい夕日"
-text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+    # Create text prompts
+    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
 
-# Create image prompt
-image = fetch_image(
-    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-)
-image_prompt = TextPrompt(
-    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-    multi_modal_data={"image": image},
-)
+    text2 = "浜辺に沈む美しい夕日"
+    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
 
-# Encode all prompts
-prompts = [text1_prompt, text2_prompt, image_prompt]
-outputs = model.encode(prompts, pooling_task="token_embed")
+    # Create image prompt
+    image = fetch_image(
+        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+    )
+    image_prompt = TextPrompt(
+        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+        multi_modal_data={"image": image},
+    )
+
+    # Encode all prompts
+    prompts = [text1_prompt, text2_prompt, image_prompt]
+    outputs = model.encode(prompts, pooling_task="token_embed")
+
+    def get_embeddings(outputs):
+        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+        embeddings = []
+        for output in outputs:
+            if VISION_START_TOKEN_ID in output.prompt_token_ids:
+                # Gather only vision tokens
+                img_start_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+                )[0][0]
+                img_end_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+                )[0][0]
+                embeddings_tensor = output.outputs.data.detach().clone()[
+                    img_start_pos : img_end_pos + 1
+                ]
+            else:
+                # Use all tokens for text-only prompts
+                embeddings_tensor = output.outputs.data.detach().clone()
+
+            # Pool and normalize embeddings
+            pooled_output = (
+                embeddings_tensor.sum(dim=0, dtype=torch.float32)
+                / embeddings_tensor.shape[0]
+            )
+            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+        return embeddings
+
+    embeddings = get_embeddings(outputs)
+
+    for embedding in embeddings:
+        print(embedding.shape)
 
 
-def get_embeddings(outputs):
-    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-    embeddings = []
-    for output in outputs:
-        if VISION_START_TOKEN_ID in output.prompt_token_ids:
-            # Gather only vision tokens
-            img_start_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-            )[0][0]
-            img_end_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-            )[0][0]
-            embeddings_tensor = output.outputs.data.detach().clone()[
-                img_start_pos : img_end_pos + 1
-            ]
-        else:
-            # Use all tokens for text-only prompts
-            embeddings_tensor = output.outputs.data.detach().clone()
-
-        # Pool and normalize embeddings
-        pooled_output = (
-            embeddings_tensor.sum(dim=0, dtype=torch.float32)
-            / embeddings_tensor.shape[0]
-        )
-        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-    return embeddings
-
-
-embeddings = get_embeddings(outputs)
-
-for embedding in embeddings:
-    print(embedding.shape)
+if __name__ == "__main__":
+    main()

examples/pooling/token_embed/multi_vector_retrieval_offline.py

@@ -4,6 +4,7 @@
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
+from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -13,6 +14,7 @@ def parse_args():
     # Set example specific arguments
     parser.set_defaults(
         model="BAAI/bge-m3",
+        pooler_config=PoolerConfig(task="token_embed"),
         runner="pooling",
         enforce_eager=True,
     )
@@ -32,15 +34,6 @@ def main(args: Namespace):
     # You should pass runner="pooling" for embedding models
     llm = LLM(**vars(args))
 
-    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
-    outputs = llm.embed(prompts)
-
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for prompt, output in zip(prompts, outputs):
-        embeds = output.outputs.embedding
-        print(len(embeds))
-
     # Generate embedding for each token. The output is a list of PoolingRequestOutput.
     outputs = llm.encode(prompts, pooling_task="token_embed")
 
@@ -50,6 +43,20 @@ def main(args: Namespace):
         multi_vector = output.outputs.data
         print(multi_vector.shape)
 
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    # Generate scores.
+    outputs = llm.score(query, documents)
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for document, output in zip(documents, outputs):
+        score = output.outputs.score
+        print(f"Pair: {[query, document]!r} \nScore: {score}")
+        print("-" * 60)
+
 
 if __name__ == "__main__":
     args = parse_args()

examples/pooling/token_embed/multi_vector_retrieval_online.py

@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
 
-vllm serve BAAI/bge-m3
+vllm serve BAAI/bge-m3 --pooler-config.task token_embed
 """
 
 import argparse
+import pprint
 
 import requests
 import torch
@@ -32,7 +33,8 @@ def parse_args():
 
 
 def main(args):
-    api_url = f"http://{args.host}:{args.port}/pooling"
+    pooling_url = f"http://{args.host}:{args.port}/pooling"
+    score_url = f"http://{args.host}:{args.port}/score"
     model_name = args.model
 
     prompts = [
@@ -43,11 +45,23 @@ def main(args):
     ]
     prompt = {"model": model_name, "input": prompts}
 
-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
     for output in pooling_response.json()["data"]:
         multi_vector = torch.tensor(output["data"])
         print(multi_vector.shape)
 
+    queries = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    prompt = {"model": model_name, "queries": queries, "documents": documents}
+    score_response = post_http_request(prompt=prompt, api_url=score_url)
+    print("\nPrompt when queries is string and documents is a list:")
+    pprint.pprint(prompt)
+    print("\nScore Response:")
+    pprint.pprint(score_response.json())
+
 
 if __name__ == "__main__":
     args = parse_args()

requirements/common.txt

@@ -32,9 +32,7 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
 mistral_common[image] >= 1.11.0
-av  # required for audio in video IO
 opencv-python-headless >= 4.13.0    # required for video IO
-soundfile  # required for audio IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.7
-flashinfer-cubin==0.6.7
+flashinfer-python==0.6.8.post1
+flashinfer-cubin==0.6.8.post1
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0

requirements/test/rocm.txt

@@ -76,9 +76,7 @@ attrs==26.1.0
 audioread==3.0.1
     # via librosa
 av==16.1.0
-    # via
-    #   -r requirements/test/../common.txt
-    #   -r requirements/test/rocm.in
+    # via -r requirements/test/rocm.in
 azure-core==1.39.0
     # via
     #   azure-identity
@@ -1333,7 +1331,6 @@ sortedcontainers==2.4.0
     # via hypothesis
 soundfile==0.13.1
     # via
-    #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
     #   genai-perf
     #   librosa

setup.py

@@ -1094,7 +1094,9 @@ setup(
         "instanttensor": ["instanttensor >= 0.1.5"],
         "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
+            "av",
             "scipy",
+            "soundfile",
             "mistral_common[audio]",
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility

tests/compile/passes/test_fuse_mla_dual_rms_norm.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit test for the MLADualRMSNormFusionPass.
+
+The pass fuses paired q/kv RMS norms in MLA attention into a single
+fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel.
+"""
+
+import pytest
+import torch
+
+import vllm.config
+from tests.compile.backend import TestBackend
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
+from vllm.config import (
+    CompilationConfig,
+    CompilationMode,
+    ModelConfig,
+    PassConfig,
+    VllmConfig,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+
+# MLA attention geometry for DeepSeek-V3 / Kimi-K2
+Q_DIM = 1536
+KV_C_DIM = 512
+K_PE_DIM = 64
+EPS = 1e-6
+
+
+class MLADualRMSNormTestModel(torch.nn.Module):
+    """
+    Minimal model reproducing the MLA dual RMS norm pattern:
+        linear -> split([q_dim, kv_dim])
+            +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear
+            +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
+                    +-- kv_c (getitem 0) -> rms_norm(kv_w, eps)
+                    +-- k_pe
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        q_dim: int = Q_DIM,
+        kv_c_dim: int = KV_C_DIM,
+        k_pe_dim: int = K_PE_DIM,
+        eps: float = EPS,
+    ):
+        super().__init__()
+        self.q_dim = q_dim
+        self.kv_dim = kv_c_dim + k_pe_dim
+        self.kv_c_dim = kv_c_dim
+        self.k_pe_dim = k_pe_dim
+
+        self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False)
+        self.q_norm = RMSNorm(q_dim, eps=eps)
+        self.kv_norm = RMSNorm(kv_c_dim, eps=eps)
+        self.q_b_proj = torch.nn.Linear(q_dim, hidden_size, bias=False)
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Avoid graph input being a direct arg to a matched pattern node
+        x = torch.relu(x)
+
+        projected = self.proj(x)
+
+        q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1)
+        kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1)
+
+        q_normed = self.q_norm(q_c)
+        kv_normed = self.kv_norm(kv_c)
+
+        q_out = self.q_b_proj(q_normed)
+        return q_out, kv_normed, k_pe
+
+    def ops_in_model_before(self):
+        return [torch.ops.vllm_ir.rms_norm.default]
+
+    def ops_in_model_after(self):
+        return [torch.ops.vllm.fused_mla_dual_rms_norm.default]
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("hidden_size", [7168])
+@pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="Only test on ROCm with AITER installed and supported",
+)
+def test_fuse_mla_dual_rms_norm(
+    dtype: torch.dtype,
+    hidden_size: int,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    torch._dynamo.reset()
+
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=dtype),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rms_norm"],
+            pass_config=PassConfig(
+                fuse_mla_dual_rms_norm=True,
+                eliminate_noops=True,
+            ),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
+        from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
+            MLADualRMSNormFusionPass,
+        )
+
+        torch.set_default_device("cuda")
+        torch.set_default_dtype(dtype)
+        torch.manual_seed(42)
+
+        m.setenv("VLLM_ROCM_USE_AITER", "1")
+        rocm_aiter_ops.refresh_env_variables()
+
+        fusion_pass = MLADualRMSNormFusionPass(vllm_config)
+        passes = [
+            NoOpEliminationPass(vllm_config),
+            fusion_pass,
+            PostCleanupPass(vllm_config),
+        ]
+        backend = TestBackend(*passes)
+        model = MLADualRMSNormTestModel(hidden_size)
+
+        x = torch.randn(1, hidden_size)
+        torch._dynamo.mark_dynamic(x, 0)
+
+        outputs_unfused = model(x)
+
+        model_fused = torch.compile(model, backend=backend)
+        outputs_fused = model_fused(x)
+
+        torch.testing.assert_close(outputs_unfused, outputs_fused, atol=1e-2, rtol=1e-2)
+
+        assert fusion_pass.matched_count == 1, (
+            f"Expected 1 fused pair, got {fusion_pass.matched_count}"
+        )
+
+        backend.check_before_ops(model.ops_in_model_before())
+        backend.check_after_ops(model.ops_in_model_after())

tests/distributed/test_eplb_events.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+import time
+
+import torch
+
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
+
+
+def test_wait_blocks_until_record():
+    event = CpuGpuEvent()
+    record_stream = torch.cuda.Stream()
+    wait_stream = torch.cuda.Stream()
+    wait_returned = threading.Event()
+
+    def waiter():
+        event.wait(stream=wait_stream)
+        wait_returned.set()
+
+    t = threading.Thread(target=waiter)
+    t.start()
+
+    time.sleep(0.05)
+    assert not wait_returned.is_set(), "wait() returned before record() was called"
+
+    event.record(stream=record_stream)
+    t.join(timeout=5.0)
+
+    assert not event._recorded.is_set()
+
+
+def test_reuse_across_multiple_cycles():
+    wrapper = CpuGpuEvent()
+    record_stream = torch.cuda.Stream()
+    wait_stream = torch.cuda.Stream()
+    NUM_CYCLES = 8
+    completed_cycles = []
+    barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)]
+
+    def waiter():
+        for i in range(NUM_CYCLES):
+            wrapper.wait(stream=wait_stream)
+            completed_cycles.append(True)
+            barriers[i].wait()
+
+    t = threading.Thread(target=waiter)
+    t.start()
+
+    for i in range(NUM_CYCLES):
+        wrapper.record(stream=record_stream)
+        barriers[i].wait()
+
+    t.join(timeout=10.0)
+    assert len(completed_cycles) == NUM_CYCLES
+
+
+def test_producer_consumer():
+    """
+    This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU
+    tensor on multiple CPU threads.
+    """
+    worker_stream = torch.cuda.Stream()
+    # Create a single element counter that will be shared between two threads
+    buf = torch.zeros(1, device="cuda")
+    NUM_ROUNDS = 5
+
+    ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)]
+    events = [CpuGpuEvent() for _ in range(NUM_ROUNDS)]
+    errors: list[str] = []
+
+    # For each round, the worker thread (writer) sets the counter in buf and waits for
+    # the main thread to read it.
+    def worker():
+        for i in range(NUM_ROUNDS):
+            if i > 0:
+                events[i - 1].wait(stream=worker_stream)
+
+            with torch.cuda.stream(worker_stream):
+                buf.fill_(float(i + 1))
+
+            worker_stream.synchronize()
+            ready_cpu[i].set()
+
+    t = threading.Thread(target=worker)
+    t.start()
+
+    for i in range(NUM_ROUNDS):
+        ready_cpu[i].wait()
+        snapshot = buf.clone()
+        events[i].record()
+        val = snapshot.item()
+        if val != float(i + 1):
+            errors.append(f"round {i}: expected {i + 1:.1f}, got {val:.1f}")
+
+    t.join(timeout=10.0)
+    assert not errors, f"Buffer ordering errors: {errors}"

tests/distributed/test_eplb_execute.py

@@ -9,7 +9,10 @@ import torch
 import torch.distributed
 
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
+from vllm.distributed.eplb.eplb_communicator import (
+    create_eplb_communicator,
+    has_nixl,
+)
 from vllm.distributed.eplb.rebalance_execute import (
     move_from_buffer,
     rearrange_expert_weights_inplace,
@@ -358,7 +361,7 @@ def _test_async_transfer_layer_without_mtp_worker(
         communicator.set_stream(cuda_stream)
 
         for layer_idx in range(num_layers):
-            is_unchanged, is_received_locally, recv_metadata = asyncio.run(
+            transfer_metadata = asyncio.run(
                 transfer_layer(
                     old_layer_indices=old_indices_cpu[layer_idx],
                     new_layer_indices=new_indices_cpu[layer_idx],
@@ -373,9 +376,7 @@ def _test_async_transfer_layer_without_mtp_worker(
             move_from_buffer(
                 expert_weights=expert_weights[layer_idx],
                 expert_weights_buffers=expert_buffer,
-                is_unchanged=is_unchanged,
-                is_received_locally=is_received_locally,
-                recv_metadata=recv_metadata,
+                transfer_metadata=transfer_metadata,
                 new_indices=new_indices_cpu[layer_idx].numpy(),
                 ep_rank=ep_rank,
             )
@@ -527,7 +528,9 @@ def _test_rearrange_expert_weights_with_redundancy(
         (4, 8, 8, 16),
     ],
 )
-@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
+@pytest.mark.parametrize(
+    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
+)
 def test_rearrange_expert_weights_with_redundancy(
     world_size,
     num_layers,
@@ -537,6 +540,8 @@ def test_rearrange_expert_weights_with_redundancy(
 ):
     """Test the functionality of rearranging expert weights with redundancy."""
 
+    if eplb_communicator == "nixl" and not has_nixl():
+        pytest.skip("NIXL is not available")
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
     distributed_run(
@@ -633,7 +638,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
         (2, 2, 2, 3),
     ],
 )
-@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
+@pytest.mark.parametrize(
+    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
+)
 def test_async_transfer_layer_without_mtp(
     world_size: int,
     num_layers: int,
@@ -643,6 +650,8 @@ def test_async_transfer_layer_without_mtp(
 ):
     """Exercise async EPLB transfer path without MTP/spec decode."""
 
+    if eplb_communicator == "nixl" and not has_nixl():
+        pytest.skip("NIXL is not available")
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
 

tests/distributed/test_eplb_utils.py

@@ -80,7 +80,7 @@ def test_commit_eplb_maps_for_layer_logical_padding():
         .contiguous()
     )
     layer = 0
-    _commit_eplb_maps_for_layer(model_state, new_phy2log, layer)
+    _commit_eplb_maps_for_layer(model_state, new_phy2log[layer], layer)
 
     assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
 
@@ -143,7 +143,7 @@ def test_commit_eplb_maps_for_layer():
     )
     new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
 
-    _commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
+    _commit_eplb_maps_for_layer(model_state, new_phy2log[0], layer=0)
 
     # Layer 0 updated
     assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])

tests/distributed/test_torchrun_example_moe.py

@@ -38,6 +38,8 @@ llm = LLM(
     distributed_executor_backend="external_launcher",
     gpu_memory_utilization=random.uniform(0.7, 0.9),
     seed=0,
+    max_model_len=1024,
+    max_num_seqs=16,
 )
 
 outputs = llm.generate(prompts, sampling_params)

tests/entrypoints/llm/test_mm_cache_external_injection.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects
+pre-processed mm_kwargs into the processor cache and reports MM cache
+hit rate metrics accurately.
+
+This is used by frameworks like Dynamo that run the HF processor on a
+frontend and transfer pre-processed mm_kwargs to the backend, avoiding
+redundant processing.
+"""
+
+import logging
+
+import pytest
+import regex as re
+
+from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
+from vllm import LLM, SamplingParams
+from vllm.renderers.params import ChatParams
+from vllm.v1.metrics import loggers as stat_loggers
+from vllm.v1.metrics.reader import Counter, Metric
+
+
+def _make_messages(image_url: str):
+    return [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": image_url},
+                },
+            ],
+        }
+    ]
+
+
+def _get_counter_value(metrics: list[Metric], name: str):
+    metric = next(m for m in metrics if m.name == name)
+    assert isinstance(metric, Counter)
+    return metric.value
+
+
+def _get_mm_cache_stats(metrics: list[Metric]):
+    mm_cache_queries = _get_counter_value(metrics, "vllm:mm_cache_queries")
+    mm_cache_hits = _get_counter_value(metrics, "vllm:mm_cache_hits")
+    return mm_cache_queries, mm_cache_hits
+
+
+def _get_mm_cache_log(llm: LLM, caplog_vllm: pytest.LogCaptureFixture) -> float:
+    caplog_vllm.clear()
+    with caplog_vllm.at_level(logging.INFO, logger=stat_loggers.__name__):
+        llm.llm_engine.do_log_stats()
+
+    assert len(caplog_vllm.records) == 1
+    msg = caplog_vllm.records[0].getMessage()
+
+    assert "MM cache hit rate" in msg
+    match = re.search(r"MM cache hit rate: ([0-9.]+)%", msg)
+    assert match is not None
+    return float(match.group(1))
+
+
+@pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:2]], indirect=True)
+@pytest.mark.parametrize("mm_processor_cache_type", ["lru", "shm"])
+def test_inject_into_mm_cache(
+    num_gpus_available,
+    image_urls,
+    mm_processor_cache_type,
+    caplog_vllm,
+):
+    """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into
+    the processor cache and MM cache hit metrics are updated correctly.
+
+    Steps:
+    1. Two normal requests (same image) -> cache miss then hit (baseline)
+    2. Extract cached kwargs, call inject_into_mm_cache with a new hash,
+       then generate with a pre-rendered input -> verifies injection works
+    """
+    llm = LLM(
+        model="llava-hf/llava-1.5-7b-hf",
+        max_model_len=4096,
+        max_num_seqs=5,
+        enforce_eager=True,
+        disable_log_stats=False,
+        limit_mm_per_prompt={"image": 2},
+        mm_processor_cache_type=mm_processor_cache_type,
+    )
+
+    # Step 1: Normal requests to populate the cache
+    llm.chat(_make_messages(image_urls[0]))
+    assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0)
+
+    llm.chat(_make_messages(image_urls[0]))
+    assert _get_mm_cache_stats(llm.get_metrics()) == (2, 1)
+    assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(50.0)
+
+    # Step 2: Use a second image to get valid expanded tokens and
+    # placeholder positions via the renderer.
+    llm.chat(_make_messages(image_urls[1]))
+    queries_before = _get_mm_cache_stats(llm.get_metrics())[0]  # 3
+
+    renderer = llm.llm_engine.renderer
+    cache = renderer.mm_processor_cache
+    assert cache is not None, "Processor cache should be enabled"
+
+    _, eng_prompts = renderer.render_chat(
+        [_make_messages(image_urls[1])],
+        ChatParams(),
+    )
+    eng_input = eng_prompts[0]
+
+    # Inject pre-processed mm_kwargs with a NEW hash via public API
+    new_mm_hash = "deadbeef" * 8
+    mm_hashes = {"image": [new_mm_hash]}
+    mm_kwargs = eng_input["mm_kwargs"]
+
+    llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)
+
+    # Build pre-rendered input (no externally_processed flag needed)
+    pre_rendered_input = {
+        "type": "multimodal",
+        "prompt_token_ids": eng_input["prompt_token_ids"],
+        "mm_kwargs": mm_kwargs,
+        "mm_hashes": mm_hashes,
+        "mm_placeholders": eng_input["mm_placeholders"],
+    }
+
+    llm.generate(
+        pre_rendered_input,
+        sampling_params=SamplingParams(max_tokens=1),
+    )
+
+    # Verify cache was queried and injection happened
+    queries_after = _get_mm_cache_stats(llm.get_metrics())[0]
+    assert queries_after > queries_before, (
+        "Cache should have been queried for the injected item"
+    )
+    mm_rate = _get_mm_cache_log(llm, caplog_vllm)
+    assert mm_rate >= 0.0, "MM cache hit rate should be reported"
+
+
+@pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:1]], indirect=True)
+def test_inject_into_mm_cache_without_cache(
+    num_gpus_available,
+    image_urls,
+):
+    """Test that inject_into_mm_cache works gracefully when processor cache
+    is disabled (mm_processor_cache_gb=0). Should not crash.
+    """
+    llm = LLM(
+        model="llava-hf/llava-1.5-7b-hf",
+        max_model_len=4096,
+        max_num_seqs=5,
+        enforce_eager=True,
+        disable_log_stats=False,
+        limit_mm_per_prompt={"image": 2},
+        mm_processor_cache_gb=0,
+    )
+
+    # Run a normal chat request first to warm up the model.
+    llm.chat(_make_messages(image_urls[0]))
+
+    # Use the renderer to get a proper EngineInput with expanded tokens
+    renderer = llm.llm_engine.renderer
+    _, eng_prompts = renderer.render_chat(
+        [_make_messages(image_urls[0])],
+        ChatParams(),
+    )
+    eng_input = eng_prompts[0]
+
+    mm_hashes = {"image": ["abcd1234" * 8]}
+    mm_kwargs = eng_input["mm_kwargs"]
+
+    # inject_into_mm_cache should not crash even without cache
+    llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs)
+
+    # Build and generate with pre-rendered input
+    pre_rendered_input = {
+        "type": "multimodal",
+        "prompt_token_ids": eng_input["prompt_token_ids"],
+        "mm_kwargs": mm_kwargs,
+        "mm_hashes": mm_hashes,
+        "mm_placeholders": eng_input["mm_placeholders"],
+    }
+
+    result = llm.generate(
+        pre_rendered_input,
+        sampling_params=SamplingParams(max_tokens=1),
+    )
+    assert len(result) == 1, "Should produce one output"
+    assert len(result[0].outputs) >= 1, "Should have at least one output sequence"

tests/entrypoints/llm/test_mm_processor_kwargs.py

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from vllm import LLM, SamplingParams
+
+
+def _make_mock_llm() -> LLM:
+    llm = object.__new__(LLM)
+    llm.model_config = SimpleNamespace(runner_type="generate")
+    return llm
+
+
+def test_generate_forwards_mm_processor_kwargs() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"num_crops": 4}
+    sampling_params = SamplingParams(max_tokens=1)
+
+    llm._run_completion = Mock(return_value=["ok"])
+
+    outputs = llm.generate(
+        "prompt",
+        sampling_params=sampling_params,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert outputs == ["ok"]
+    assert llm._run_completion.call_args.kwargs["mm_processor_kwargs"] == (
+        mm_processor_kwargs
+    )
+
+
+def test_enqueue_forwards_mm_processor_kwargs() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"do_resize": False}
+    sampling_params = SamplingParams(max_tokens=1)
+
+    llm._add_completion_requests = Mock(return_value=["req-0"])
+
+    request_ids = llm.enqueue(
+        "prompt",
+        sampling_params=sampling_params,
+        use_tqdm=False,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert request_ids == ["req-0"]
+    assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (
+        mm_processor_kwargs
+    )
+
+
+def test_chat_forwards_mm_processor_kwargs() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"do_pan_and_scan": True}
+    sampling_params = SamplingParams(max_tokens=1)
+    messages = [{"role": "user", "content": "hello"}]
+
+    llm._run_chat = Mock(return_value=["ok"])
+
+    outputs = llm.chat(
+        messages,
+        sampling_params=sampling_params,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert outputs == ["ok"]
+    assert llm._run_chat.call_args.kwargs["mm_processor_kwargs"] == (
+        mm_processor_kwargs
+    )
+
+
+def test_run_completion_forwards_mm_processor_kwargs() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"min_pixels": 4 * 28 * 28}
+    sampling_params = SamplingParams(max_tokens=1)
+    sentinel_output = ["done"]
+
+    llm._add_completion_requests = Mock()
+    llm._run_engine = Mock(return_value=sentinel_output)
+
+    outputs = llm._run_completion(
+        prompts=["prompt"],
+        params=sampling_params,
+        output_type=object,
+        use_tqdm=False,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert outputs == sentinel_output
+    assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (
+        mm_processor_kwargs
+    )
+
+
+def test_add_completion_requests_forwards_mm_processor_kwargs() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"max_dynamic_patch": 4}
+    sampling_params = SamplingParams(max_tokens=1)
+
+    llm._params_to_seq = Mock(return_value=[sampling_params])
+    llm._lora_request_to_seq = Mock(return_value=[None])
+    llm._priority_to_seq = Mock(return_value=[0])
+    llm._preprocess_cmpl_one = Mock(return_value={"prompt_token_ids": [1]})
+
+    captured_prompts = []
+
+    def fake_render_and_add_requests(*, prompts, **_kwargs):
+        captured_prompts.extend(prompts)
+        return ["req-0"]
+
+    llm._render_and_add_requests = Mock(side_effect=fake_render_and_add_requests)
+
+    request_ids = llm._add_completion_requests(
+        prompts=["prompt"],
+        params=sampling_params,
+        use_tqdm=False,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert request_ids == ["req-0"]
+    llm._preprocess_cmpl_one.assert_called_once_with(
+        "prompt",
+        None,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+    assert captured_prompts == [{"prompt_token_ids": [1]}]
+
+
+def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"num_crops": 8}
+    prompt = {"prompt": "<image>", "multi_modal_data": {"image": object()}}
+
+    renderer = Mock()
+    renderer.default_cmpl_tok_params = Mock()
+    renderer.default_cmpl_tok_params.with_kwargs.return_value = "tok-params"
+    renderer.render_cmpl.return_value = ["engine-input"]
+    llm.renderer = renderer
+
+    monkeypatch.setattr(
+        "vllm.entrypoints.llm.parse_model_prompt",
+        lambda _model_config, parsed_prompt: parsed_prompt,
+    )
+
+    outputs = llm._preprocess_cmpl(
+        [prompt],
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert outputs == ["engine-input"]
+    renderer.render_cmpl.assert_called_once_with(
+        [prompt],
+        "tok-params",
+        prompt_extras={"mm_processor_kwargs": mm_processor_kwargs},
+    )
+
+
+def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    llm = _make_mock_llm()
+    prompt = {
+        "prompt": "<image>",
+        "multi_modal_data": {"image": object()},
+        "mm_processor_kwargs": {"num_crops": 2},
+    }
+
+    renderer = Mock()
+    renderer.default_cmpl_tok_params = Mock()
+    renderer.default_cmpl_tok_params.with_kwargs.return_value = "tok-params"
+    renderer.render_cmpl.return_value = ["engine-input"]
+    llm.renderer = renderer
+
+    monkeypatch.setattr(
+        "vllm.entrypoints.llm.parse_model_prompt",
+        lambda _model_config, parsed_prompt: parsed_prompt,
+    )
+
+    outputs = llm._preprocess_cmpl([prompt])
+
+    assert outputs == ["engine-input"]
+    renderer.render_cmpl.assert_called_once_with(
+        [prompt],
+        "tok-params",
+        prompt_extras=None,
+    )
+
+
+def test_preprocess_chat_applies_mm_processor_kwargs_to_renderer() -> None:
+    llm = _make_mock_llm()
+    mm_processor_kwargs = {"num_crops": 8}
+    messages = [[{"role": "user", "content": "Describe this image."}]]
+
+    renderer = Mock()
+    renderer.tokenizer = object()
+    renderer.default_chat_tok_params = Mock()
+    renderer.default_chat_tok_params.with_kwargs.return_value = "tok-params"
+    renderer.render_chat.return_value = (messages, ["engine-input"])
+    llm.renderer = renderer
+
+    outputs = llm._preprocess_chat(
+        messages,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+
+    assert outputs == ["engine-input"]
+    call_args = renderer.render_chat.call_args
+    assert call_args.args[0] == messages
+    assert call_args.args[1].mm_processor_kwargs == mm_processor_kwargs
+    assert call_args.args[2] == "tok-params"
+    assert call_args.kwargs["prompt_extras"] == {
+        "mm_processor_kwargs": mm_processor_kwargs
+    }
+
+
+def test_preprocess_chat_omits_mm_processor_kwargs_when_no_override() -> None:
+    llm = _make_mock_llm()
+    messages = [[{"role": "user", "content": "Describe this image."}]]
+
+    renderer = Mock()
+    renderer.tokenizer = object()
+    renderer.default_chat_tok_params = Mock()
+    renderer.default_chat_tok_params.with_kwargs.return_value = "tok-params"
+    renderer.render_chat.return_value = (messages, ["engine-input"])
+    llm.renderer = renderer
+
+    outputs = llm._preprocess_chat(messages)
+
+    assert outputs == ["engine-input"]
+    call_args = renderer.render_chat.call_args
+    assert call_args.args[0] == messages
+    assert call_args.args[1].mm_processor_kwargs is None
+    assert call_args.args[2] == "tok-params"
+    assert call_args.kwargs["prompt_extras"] is None

tests/entrypoints/pooling/classify/test_offline.py

@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
 import torch
 
 from tests.models.utils import softmax
-from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
+from vllm import LLM, ClassificationRequestOutput, PoolingParams
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask
 
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.probs) == num_labels
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_classify(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    assert len(outputs) == 1
-    assert isinstance(outputs[0], PoolingRequestOutput)
-    assert outputs[0].prompt_token_ids == prompt_token_ids
-    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
         llm.score("ping", "pong", use_tqdm=False)
 
 
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Embedding API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/classify/test_online.py

@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
-    task = "token_classify"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-    poolings = PoolingResponse.model_validate(response.json())
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 8
-    assert len(poolings.data[0].data[0]) == 2
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
+
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -38,11 +37,11 @@ def llm():
         seed=0,
         attention_config=attention_config,
     )
+    assert embedding_size == llm.model_config.embedding_size
 
     yield weakref.proxy(llm)
 
     del llm
-
     cleanup_dist_env_and_memory()
 
 
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.embedding) == embedding_size
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_embed(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    multi_vector = outputs[0].outputs.data
-    assert multi_vector.shape == (11, 384)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
     )
 
 
-@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["token_classify", "classify", "token_embed", "plugin"]
+)
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Classification API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/embed/test_online.py

@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
-    task = "token_embed"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 384
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["classify", "token_classify", "token_embed", "plugin"]
+)
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/scoring/test_cross_encoder_online.py

@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
     assert len(poolings.data[0].data) == 1
 
 
-@pytest.mark.asyncio
-async def test_pooling_token_classify(server: RemoteOpenAIServer):
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": MODEL_NAME,
-            "task": "token_classify",
-            "input": input_text,
-            "encoding_format": "float",
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 1
-
-
 @pytest.mark.asyncio
 async def test_rerank_max_tokens_per_doc(
     server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     response = requests.post(
         server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_classify/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "Scoring API is only enabled for num_labels == 1."
+    err_msg = "This model does not support the Scoring API."
     with pytest.raises(ValueError, match=err_msg):
         llm.score("ping", "pong", use_tqdm=False)
 
 
 @pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "classify":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Embedding API is not supported by this model.+"
+        err_msg = "Embedding API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_classify/test_online.py

@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "embed":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Classification API is not supported by this model.+"
+        err_msg = "Classification API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_embed/test_online.py

@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/ir/ir_test_utils.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shared test utilities for vLLM IR op correctness tests.
+"""
+
+import torch
+
+from vllm.ir.op import IrOp
+
+NUM_TOKENS = [1, 8, 17, 32, 512, 2048]
+COMMON_HIDDEN_SIZES = [
+    2048,  # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
+    4096,  # Llama 3 8B, Qwen 3 8B
+    5120,  # Llama 4 Scout 17B-16E
+    7168,  # DeepSeek V3
+    8192,  # Llama 3 70B
+]
+
+
+def clone_args(args: tuple) -> tuple:
+    return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
+
+
+def supported_providers(op: IrOp) -> list[str]:
+    return [
+        name for name, impl in op.impls.items() if name != "native" and impl.supported
+    ]
+
+
+def assert_close(op: IrOp, actual, expected):
+    if isinstance(actual, torch.Tensor):
+        tol = op.get_tolerance(actual.dtype)
+        try:
+            torch.testing.assert_close(actual, expected, **tol)
+        except AssertionError as e:
+            raise AssertionError(
+                f"{e}\n\nTo adjust tolerance, use:\n"
+                f"  ir.ops.{op.name}.override_tolerance("
+                f"{actual.dtype}, atol=..., rtol=...)"
+            ) from None
+    elif isinstance(actual, (tuple, list)):
+        for a, ex in zip(actual, expected):
+            assert_close(op, a, ex)
+    else:
+        assert actual == expected

tests/ir/test_op.py

@@ -495,3 +495,68 @@ def test_uuid_and_oot(tmp_path: Path):
     assert uuid2 == uuid
     assert uuid2 != uuid1
     del _custom_mm.impls["impl_mm_oot"]
+
+
+def _test_native(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return x + y
+
+
+def _make_op_with_generator(name: str = "_ig_test"):
+    op = IrOp(name, _test_native)
+
+    @op.register_input_generator
+    def _gen(n: int = 4):
+        x = torch.randn(n, 3)
+        y = torch.randn(n, 3)
+        return x, y
+
+    return op
+
+
+def _test_native_single(x: torch.Tensor) -> torch.Tensor:
+    return x
+
+
+class TestInputGenerator:
+    def test_no_input_generator_by_default(self):
+        op = IrOp("_ig_test_no_gen", _test_native_single)
+        assert not op.has_input_generator
+
+    def test_register_input_generator(self):
+        op = _make_op_with_generator("_ig_test_reg")
+        assert op.has_input_generator
+
+    def test_generate_inputs_returns_tuple(self):
+        op = _make_op_with_generator("_ig_test_tuple")
+        result = op.generate_inputs(n=2)
+        assert isinstance(result, tuple)
+        assert len(result) == 2
+        assert result[0].shape == (2, 3)
+        assert result[1].shape == (2, 3)
+
+    def test_generate_inputs_default_kwargs(self):
+        op = _make_op_with_generator("_ig_test_default")
+        result = op.generate_inputs()
+        assert result[0].shape == (4, 3)
+
+    def test_generate_inputs_without_registration_raises(self):
+        op = IrOp("_ig_test_no_gen_raises", _test_native_single)
+        with pytest.raises(RuntimeError, match="No input generator"):
+            op.generate_inputs()
+
+
+class TestTolerance:
+    def test_override_and_get_tolerance(self):
+        op = IrOp("_tol_test", _test_native)
+
+        tol = op.get_tolerance(torch.float32)
+        assert tol == {"atol": 1e-5, "rtol": 1.3e-6}
+
+        op.override_tolerance(torch.float32, atol=0.1, rtol=0.2)
+        assert op.get_tolerance(torch.float32) == {"atol": 0.1, "rtol": 0.2}
+        assert op.get_tolerance(torch.float16) == {"atol": 1e-3, "rtol": 1e-3}
+
+    def test_get_tolerance_raises_for_unknown_dtype(self):
+        op = IrOp("_tol_test_unknown", _test_native)
+        with pytest.raises(ValueError, match="No tolerance defined"):
+            op.get_tolerance(torch.complex64)

tests/kernels/attention/test_triton_decode_attention.py

@@ -4,9 +4,12 @@
 import pytest
 import torch
 
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @pytest.mark.parametrize("B", [3, 5])
 @pytest.mark.parametrize("L", [1027, 1025])
@@ -25,33 +28,35 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
 
     num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
     req_to_page = torch.randint(
-        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
+        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
     )
     req_to_token = req_to_page * PAGE_SIZE
     req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
-    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1)
+    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
+        1, 1, -1
+    )
     req_to_token = req_to_token.view(B, -1)
     req_to_token = req_to_token[:, :seq_len].contiguous()
 
     # q represents the new token being generated, one per batch
-    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
+    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
 
     # k_buffer and v_buffer represent all previous tokens
     # Page size is 1.
-    k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")
-    v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda")
+    k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
+    v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
 
     # o will have the same shape as q
-    o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
+    o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
 
-    lse = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
+    lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
 
-    b_seq_len = torch.full((B,), seq_len, device="cuda")
+    b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
 
     attn_logits = torch.empty(
         (B, H_Q, num_kv_splits, D_V + 1),
         dtype=torch.float32,
-        device="cuda",
+        device=DEVICE_TYPE,
     )
 
     # Call the original implementation.
@@ -127,25 +132,27 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
 
     num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
     req_to_page = torch.randint(
-        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
+        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
     )
     req_to_token = req_to_page * PAGE_SIZE
     req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
-    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1)
+    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
+        1, 1, -1
+    )
     req_to_token = req_to_token.view(B, -1)
     req_to_token = req_to_token[:, :seq_len].contiguous()
 
-    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
+    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
 
     # Create BF16 K/V as reference
-    k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")
-    v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda")
+    k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
+    v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
 
     # --- BF16 reference ---
-    o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
-    lse_ref = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
+    o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
+    lse_ref = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
     attn_logits = torch.empty(
-        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
+        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
     )
 
     if PAGE_SIZE == 1:
@@ -156,7 +163,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
             o_ref,
             lse_ref,
             req_to_token,
-            b_seq_len=torch.full((B,), seq_len, device="cuda"),
+            b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
             attn_logits=attn_logits,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -171,7 +178,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
             o_ref,
             lse_ref,
             req_to_page,
-            b_seq_len=torch.full((B,), seq_len, device="cuda"),
+            b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
             attn_logits=attn_logits,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -182,10 +189,10 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
     k_fp8, k_scale = _quantize_to_fp8(k_bf16)
     v_fp8, v_scale = _quantize_to_fp8(v_bf16)
 
-    o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
-    lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
+    o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
+    lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
     attn_logits_fp8 = torch.empty(
-        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
+        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
     )
 
     if PAGE_SIZE == 1:
@@ -196,7 +203,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
             o_fp8,
             lse_fp8,
             req_to_token,
-            b_seq_len=torch.full((B,), seq_len, device="cuda"),
+            b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
             attn_logits=attn_logits_fp8,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -213,7 +220,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
             o_fp8,
             lse_fp8,
             req_to_page,
-            b_seq_len=torch.full((B,), seq_len, device="cuda"),
+            b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
             attn_logits=attn_logits_fp8,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,

tests/kernels/attention/test_triton_prefill_attention.py

@@ -5,8 +5,11 @@ import pytest
 import torch
 import torch.nn.functional as F
 
+from vllm.platforms import current_platform
 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def ref_masked_attention(
     q: torch.Tensor,
@@ -92,17 +95,19 @@ def test_context_attention(
     torch.manual_seed(42)
 
     # Generate random sequence lengths for each batch
-    seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
+    seq_lens = torch.randint(
+        max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
+    )
     total_tokens = seq_lens.sum().item()
 
     # Create batch start locations
-    b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
+    b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
     b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
 
     # Create input tensors
-    q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
-    k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
-    v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
+    q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
+    k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
+    v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
     o = torch.zeros_like(q)
 
     # Call Triton kernel
@@ -169,17 +174,19 @@ def test_context_attention_sliding_window(
     torch.manual_seed(42)
 
     # Generate random sequence lengths for each batch
-    seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
+    seq_lens = torch.randint(
+        max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
+    )
     total_tokens = seq_lens.sum().item()
 
     # Create batch start locations
-    b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
+    b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
     b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
 
     # Create input tensors
-    q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
-    k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
-    v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
+    q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
+    k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
+    v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
     o = torch.zeros_like(q)
 
     # Call Triton kernel

tests/kernels/attention/test_triton_unified_attention.py

@@ -10,6 +10,8 @@ from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
 
+DEVICE_TYPE = current_platform.device_type
+
 NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
 HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]
@@ -114,7 +116,7 @@ def test_triton_unified_attn(
     q_dtype: torch.dtype | None,
     seq_threshold_3D: int,
 ) -> None:
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     set_random_seed(0)
     num_seqs = len(seq_lens)
@@ -249,7 +251,7 @@ def test_triton_unified_attn_fp16_input_fp8_output(
     seq_threshold_3D: int,
 ) -> None:
     """Test with fp16 input and fp8 output using output_scale."""
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     set_random_seed(0)
     num_seqs = len(seq_lens)

tests/kernels/attention/test_use_trtllm_attention.py

@@ -62,7 +62,7 @@ def test_supports_batch_invariant_disables():
 
 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
 @patch(
-    "vllm.utils.flashinfer.current_platform.is_device_capability",
+    "vllm.utils.flashinfer.current_platform.is_device_capability_family",
     return_value=True,
 )
 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
@@ -72,7 +72,7 @@ def test_supports_sm100_with_artifactory(_art, _cap):
 
 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
 @patch(
-    "vllm.utils.flashinfer.current_platform.is_device_capability",
+    "vllm.utils.flashinfer.current_platform.is_device_capability_family",
     return_value=False,
 )
 def test_supports_non_sm100_platform(_cap):
@@ -81,7 +81,7 @@ def test_supports_non_sm100_platform(_cap):
 
 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
 @patch(
-    "vllm.utils.flashinfer.current_platform.is_device_capability",
+    "vllm.utils.flashinfer.current_platform.is_device_capability_family",
     return_value=True,
 )
 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)

tests/kernels/ir/test_ir_ops.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Meta-tests for vLLM IR op infrastructure.
+
+Ensures all registered ops have input generators defined.
+Per-op correctness tests live alongside their op definitions
+(e.g. tests/kernels/ir/test_layernorm.py).
+"""
+
+import vllm.kernels  # noqa: F401 — registers provider implementations
+from vllm.ir.op import IrOp
+
+
+def test_all_ops_have_input_generator():
+    missing = [name for name, op in IrOp.registry.items() if not op.has_input_generator]
+    assert not missing, (
+        f"IR ops without input generators: {missing}. "
+        f"Register one with @ir.ops.<name>.register_input_generator"
+    )

tests/kernels/ir/test_layernorm.py

@@ -5,17 +5,17 @@ import torch
 
 # This registers op implementations
 import vllm.kernels  # noqa: F401
+from tests.ir.ir_test_utils import (
+    COMMON_HIDDEN_SIZES,
+    NUM_TOKENS,
+    assert_close,
+    clone_args,
+    supported_providers,
+)
 from tests.kernels.allclose_default import get_default_rtol
 from vllm import ir
 from vllm.platforms import current_platform
 
-
-def rms_norm_inputs(n_tokens: int, hidden_size: int, dtype: torch.dtype):
-    x = torch.randn(n_tokens, hidden_size, dtype=dtype)
-    weight = torch.rand(hidden_size, dtype=dtype)
-    return x, weight
-
-
 rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
 
 
@@ -40,8 +40,8 @@ def test_rms_norm_registration():
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("n_tokens", [1, 8, 17])
-@pytest.mark.parametrize("hidden_size", [16, 4096, 8192])
+@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
 @pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
@@ -53,7 +53,9 @@ class TestRMSNorm:
         torch.set_default_device(current_platform.device_type)
 
     def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
-        x, weight = rms_norm_inputs(4, 8, dtype)
+        x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
+            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
+        )
         out = rms_norm_native(x, weight, epsilon=epsilon)
 
         # Check shape, dtype, device
@@ -71,59 +73,59 @@ class TestRMSNorm:
         out4 = rms_norm_native(x, None, epsilon=epsilon)
         torch.testing.assert_close(out3, out4)
 
-    @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels"])
+    @pytest.mark.parametrize("provider", supported_providers(ir.ops.rms_norm))
     def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
         impl = ir.ops.rms_norm.impls[provider]
-        if not impl.supported:
-            pytest.skip(f"{provider} impl not supported on this platform")
-
-        x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
-        args = (x, weight, epsilon, None)
-
-        assert impl.supported
-
-        if provider == "aiter" and dtype not in [torch.float16, torch.bfloat16]:
-            assert not impl.supports_args(*args)
-            return
-
-        assert impl.supports_args(*args)
-
-        out_impl = impl.impl_fn(*args)
-        out_native = rms_norm_native(*args)
-
-        torch.testing.assert_close(
-            out_impl, out_native, rtol=get_default_rtol(out_impl), atol=1e-3
+        x, weight, eps = ir.ops.rms_norm.generate_inputs(
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
         )
+        args = (x, weight, eps)
+
+        if not impl.supports_args(*args):
+            pytest.skip(f"{provider} does not support args")
+
+        ref_output = rms_norm_native(*clone_args(args))
+        output = impl.impl_fn(*clone_args(args))
+        assert_close(ir.ops.rms_norm, output, ref_output)
 
         # check that dispatched call matches direct call
         with ir.ops.rms_norm.set_priority([provider, "native"]):
-            out_impl2 = ir.ops.rms_norm(*args)
-
-        # exact match
-        torch.testing.assert_close(out_impl2, out_impl, rtol=0.0, atol=0.0)
+            out_dispatched = ir.ops.rms_norm(*args)
+        out_direct = impl.impl_fn(*args)
+        torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
 
         # none of these support variance_size override
-        assert not impl.supports_args(x, weight, epsilon, 4)
-        assert not impl.supports_args(x, weight, epsilon, variance_size=4)
+        assert not impl.supports_args(x, weight, eps, 4)
+        assert not impl.supports_args(x, weight, eps, variance_size=4)
 
         # test weight=None behavior
-        out_impl_no_weight = impl.impl_fn(x, None, epsilon)
-        out_impl_unit_weight = impl.impl_fn(x, torch.ones_like(weight), epsilon)
-        torch.testing.assert_close(
-            out_impl_no_weight,
-            out_impl_unit_weight,
-            rtol=get_default_rtol(out_impl_no_weight),
-            atol=2e-4,
-        )
+        out_no_weight = impl.impl_fn(x, None, eps)
+        out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps)
+        assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight)
 
     @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
     def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
         if not ir.ops.rms_norm.impls[provider].supported:
             pytest.skip(f"{provider} impl not supported on this platform")
 
-        x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
-        args = (x, weight, epsilon, None)
+        args = ir.ops.rms_norm.generate_inputs(
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
+        )
 
         # When checking the torch op, we have to set priority and use dispatch
         with ir.ops.rms_norm.set_priority([provider, "native"]):
             torch.library.opcheck(torch.ops.vllm_ir.rms_norm, args)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="aiter is only supported on ROCm",
+)
+def test_aiter_rejects_unsupported_dtypes():
+    torch.set_default_device(current_platform.device_type)
+    impl = ir.ops.rms_norm.impls["aiter"]
+    for dtype in [torch.float32, torch.float64]:
+        args = ir.ops.rms_norm.generate_inputs(
+            num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
+        )
+        assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"

tests/kernels/moe/test_fused_topk.py

@@ -135,3 +135,70 @@ def test_fused_topk_bias(
         topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
     )
     torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize("num_experts", [6, 8, 16])
+@pytest.mark.parametrize("topk", [3, 4])
+@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
+@pytest.mark.parametrize("bad_value", [float("nan"), float("inf")])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
+def test_fused_topk_nan_inf_clamp(
+    num_experts: int,
+    topk: int,
+    scoring_func: str,
+    bad_value: float,
+    dtype: torch.dtype,
+):
+    """Regression test for the NaN/Inf clamp in topk_softmax_kernels.cu.
+
+    Degenerate hidden states (e.g., from CUDA graph padding) can produce
+    NaN/Inf gating logits. Without the clamp, softmax/sigmoid outputs are
+    NaN and the argmax loop picks expert 0 for every top-k slot (since
+    "NaN > NaN" is false per IEEE 754), yielding duplicate expert IDs that
+    crash downstream MoE sort kernels. The fix clamps NaN/Inf to 0 before
+    argmax so index tie-breaking selects unique experts [0, 1, ..., k-1].
+    """
+    torch.manual_seed(0)
+    num_tokens = 4
+    hidden_size = 1024
+    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
+
+    # Row 0: all normal. Rows 1-3: fully poisoned with NaN or Inf.
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+    gating_output[1:, :] = bad_value
+
+    topk_weights, topk_ids, _ = fused_topk(
+        hidden_states=hidden_states,
+        gating_output=gating_output,
+        topk=topk,
+        renormalize=False,
+        scoring_func=scoring_func,
+    )
+
+    # Normal row must still match the torch reference.
+    ref_weights, ref_ids = torch_topk(
+        gating_output=gating_output[:1],
+        topk=topk,
+        renormalize=False,
+        scoring_func=scoring_func,
+    )
+    torch.testing.assert_close(
+        ref_weights.to(torch.float32), topk_weights[:1], atol=1e-2, rtol=1e-2
+    )
+    torch.testing.assert_close(ref_ids.to(torch.int32), topk_ids[:1], atol=0, rtol=0)
+
+    # Poisoned rows: IDs must be unique (no duplicates) and weights must be
+    # finite (no NaN/Inf propagation into downstream MoE kernels).
+    for row in range(1, num_tokens):
+        row_ids = topk_ids[row]
+        assert row_ids.unique().numel() == topk, (
+            f"Row {row} has duplicate expert IDs {row_ids.tolist()} "
+            f"(bad_value={bad_value}, scoring_func={scoring_func})"
+        )
+        assert torch.isfinite(topk_weights[row]).all(), (
+            f"Row {row} has non-finite weights {topk_weights[row].tolist()} "
+            f"(bad_value={bad_value}, scoring_func={scoring_func})"
+        )

tests/kernels/moe/test_moe_layer.py

@@ -236,7 +236,6 @@ class MoETestConfig:
     use_gate: bool
     use_routed_input_transform: bool
     enable_eplb: bool = False
-    reduce_results: bool = False
     backend: str | None = None
     ep_size: int = 1
     dp_size: int = 1
@@ -295,7 +294,6 @@ def generate_valid_test_configs(
         use_shared_experts,
         use_gate,
         use_routed_input_transform,
-        reduce_results,
     ) in product(
         SHAPE_COMBOS,
         NUM_EXPERTS,
@@ -304,7 +302,6 @@ def generate_valid_test_configs(
         [False, True],  # shared
         [False, True],  # gate
         [False, True],  # routed input exform
-        [False, True],  # reduce results
     ):
         config = MoETestConfig(
             shape[0],  # m
@@ -318,7 +315,6 @@ def generate_valid_test_configs(
             use_gate,
             use_routed_input_transform,
             enable_eplb,
-            reduce_results,
             backend,
             ep_size,
             dp_size,
@@ -395,18 +391,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
         and config.backend.startswith("flashinfer_nvlink")
         and not current_platform.has_device_capability(90)
     ):
-        return False, "flashinfer_nvlink needs an H100+ GPUs"
-
-    # reduce_results incompatibilities
-    if config.reduce_results and config.use_shared_experts:
-        return False, "reduce_results=True is not compatible with shared_experts=True"
-
-    if config.reduce_results and config.quantization is not None:
-        return (
-            False,
-            "reduce_results=True only tested with unquantized data types in "
-            "order to limit number of tests run",
-        )
+        return False, "flashinfer_nvlink needs H100+ GPUs"
 
     # Backend-specific checks
     if config.backend is not None:
@@ -448,10 +433,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
     if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS:
         return False, f"EPLB not supported with {config.backend}."
 
-    world_size = config.tp_size * config.dp_size
-    if config.reduce_results and world_size == 1:
-        return False, "reduce_results=True only makes sense for multi-GPU tests"
-
     if (
         config.backend is not None
         and config.backend.startswith("flashinfer_nvlink")
@@ -846,7 +827,6 @@ def make_fused_moe_layer(
     tp_size: int,
     ep_size: int,
     dp_size: int,
-    reduce_results: bool,
     w1: torch.Tensor,
     w2: torch.Tensor,
     top_k: int,
@@ -874,7 +854,7 @@ def make_fused_moe_layer(
     routed_input_transform: torch.nn.Module | None = None,
     routed_output_transform: torch.nn.Module | None = None,
     pcp_size: int | None = 1,
-) -> tuple[Callable, FusedMoE]:
+) -> FusedMoE:
     quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)
 
     kwargs = dict()
@@ -887,8 +867,10 @@ def make_fused_moe_layer(
     # Add gate and routed_input_transform if provided
     if gate is not None:
         kwargs["gate"] = gate
+
     if routed_input_transform is not None:
         kwargs["routed_input_transform"] = routed_input_transform
+        kwargs["routed_output_transform"] = routed_output_transform
 
     layer = builder(
         num_experts=global_num_experts,
@@ -896,7 +878,6 @@ def make_fused_moe_layer(
         hidden_size=hidden_size,
         intermediate_size=intermediate_size,
         params_dtype=in_dtype,
-        reduce_results=reduce_results,
         renormalize=renormalize,
         use_grouped_topk=use_grouped_topk,
         num_expert_group=num_expert_group,
@@ -936,36 +917,7 @@ def make_fused_moe_layer(
 
     layer.quant_method.process_weights_after_loading(layer)
 
-    def _moe(
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-    ) -> torch.Tensor:
-        if shared_experts is None:
-            final_shared_states = None
-            final_hidden_states = layer(hidden_states, router_logits)
-        else:
-            final_shared_states, final_hidden_states = layer(
-                hidden_states, router_logits
-            )
-
-        # Apply routed output transform if provided
-        # (e.g., latent space -> original space)
-        if routed_output_transform is not None:
-            final_hidden_states = routed_output_transform(final_hidden_states)
-
-        if shared_experts is not None:
-            assert not reduce_results
-            assert final_shared_states is not None
-            final_hidden_states += final_shared_states
-
-        if not reduce_results and layer.tp_size > 1:
-            final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel(
-                final_hidden_states
-            )
-
-        return final_hidden_states
-
-    return _moe, layer
+    return layer
 
 
 def make_fake_moe_layer(
@@ -999,7 +951,6 @@ def make_fake_moe_layer(
     tp_size: int = 1,
     dp_size: int = 1,
     ep_size: int = 1,
-    reduce_results: bool = False,
 ) -> Callable:
     activation = MoEActivation.from_str(activation)
 
@@ -1101,7 +1052,7 @@ def make_fake_moe_layer(
 
 
 def _test_body_regular(
-    moe_fn: Callable,
+    moe_layer: Callable,
     hidden_states: torch.Tensor,
     router_logits: torch.Tensor,
     vllm_config: VllmConfig,
@@ -1118,13 +1069,12 @@ def _test_body_regular(
         num_tokens=num_tokens,
         num_tokens_across_dp=num_tokens_across_dp,
     ):
-        output = moe_fn(hidden_states, router_logits)
+        output = moe_layer(hidden_states, router_logits)
 
     return baseline_output, output
 
 
 def _test_body_eplb(
-    moe_fn: Callable,
     moe_layer: FusedMoE,
     hidden_states: torch.Tensor,
     router_logits: torch.Tensor,
@@ -1145,7 +1095,6 @@ def _test_body_eplb(
     n: int,
     top_k: int,
     shared_experts,
-    reduce_results: bool,
     gate: torch.nn.Module | None,
     routed_input_transform: torch.nn.Module | None,
     routed_output_transform: torch.nn.Module | None,
@@ -1161,7 +1110,7 @@ def _test_body_eplb(
         num_tokens=num_tokens,
         num_tokens_across_dp=num_tokens_across_dp,
     ):
-        output_before = moe_fn(hidden_states, router_logits)
+        output_before = moe_layer(hidden_states, router_logits)
 
     # Create a fresh FusedMoE layer with enable_eplb=True
     # Delete the original layer's registration so the constructor can
@@ -1174,7 +1123,7 @@ def _test_body_eplb(
     # When using routed_input_transform, experts operate in latent space
     hidden_size_for_layer = k // 2 if routed_input_transform is not None else k
 
-    moe_fn, moe_layer = make_fused_moe_layer(
+    eplb_moe_layer = make_fused_moe_layer(
         quantization=quantization,
         use_ep=use_ep,
         hidden_size=hidden_size_for_layer,
@@ -1183,7 +1132,6 @@ def _test_body_eplb(
         tp_size=tp_size,
         ep_size=ep_size,
         dp_size=dp_size,
-        reduce_results=reduce_results,
         w1=w1,
         w2=w2,
         top_k=top_k,
@@ -1196,14 +1144,14 @@ def _test_body_eplb(
     )
 
     # Necessary?
-    if moe_layer._expert_map is not None:
-        moe_layer._expert_map = moe_layer._expert_map.to(device)
+    if eplb_moe_layer._expert_map is not None:
+        eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device)
 
     # All ranks must generate the same permutation
     initial_indices = torch.arange(num_experts, dtype=torch.long)
     shuffled_indices = initial_indices[torch.randperm(num_experts)]
 
-    expert_weights = [list(moe_layer.get_expert_weights())]
+    expert_weights = [list(eplb_moe_layer.get_expert_weights())]
 
     communicator = create_eplb_communicator(
         group_coordinator=get_eplb_group(),
@@ -1227,7 +1175,7 @@ def _test_body_eplb(
         num_experts, dtype=torch.int32, device=device
     )
 
-    moe_layer.set_eplb_state(
+    eplb_moe_layer.set_eplb_state(
         moe_layer_idx=0,
         expert_load_view=torch.zeros(
             (1, num_experts),
@@ -1244,7 +1192,7 @@ def _test_body_eplb(
         ),
     )
 
-    moe_layer.eplb_state.should_record_tensor = torch.ones(
+    eplb_moe_layer.eplb_state.should_record_tensor = torch.ones(
         (), dtype=torch.bool, device=device
     )
 
@@ -1255,7 +1203,7 @@ def _test_body_eplb(
         num_tokens=num_tokens,
         num_tokens_across_dp=num_tokens_across_dp,
     ):
-        output_after = moe_fn(hidden_states, router_logits)
+        output_after = eplb_moe_layer(hidden_states, router_logits)
 
     return output_before, output_after
 
@@ -1274,7 +1222,6 @@ def _run_one_config(
     num_experts: int,
     top_k: int,
     quantization: str | None,
-    reduce_results: bool,
     backend: str | None,
     test_body_fn: Callable,
     use_shared_experts: bool,
@@ -1341,7 +1288,6 @@ def _run_one_config(
         tp_size=tp_size,
         ep_size=ep_size,
         dp_size=dp_size,
-        reduce_results=reduce_results,
     )
 
     baseline_output = baseline_layer(hidden_states, router_logits)
@@ -1369,7 +1315,7 @@ def _run_one_config(
         hidden_size_for_layer = k // 2 if routed_input_transform is not None else k
 
         # Create initial MoE layer
-        moe_fn, moe_layer = make_fused_moe_layer(
+        moe_layer = make_fused_moe_layer(
             quantization=quantization,
             use_ep=use_ep,
             hidden_size=hidden_size_for_layer,
@@ -1378,7 +1324,6 @@ def _run_one_config(
             tp_size=tp_size,
             ep_size=ep_size,
             dp_size=dp_size,
-            reduce_results=reduce_results,
             w1=w1,
             w2=w2,
             top_k=top_k,
@@ -1402,7 +1347,6 @@ def _run_one_config(
 
         # Call the test body function with all necessary context
         expected, actual = test_body_fn(
-            moe_fn=moe_fn,
             moe_layer=moe_layer,
             hidden_states=hidden_states,
             router_logits=router_logits,
@@ -1423,7 +1367,6 @@ def _run_one_config(
             m=m,
             top_k=top_k,
             shared_experts=shared_experts,
-            reduce_results=reduce_results,
             gate=gate,
             routed_input_transform=routed_input_transform,
             routed_output_transform=routed_output_transform,
@@ -1520,7 +1463,6 @@ def test_moe_layer_no_parallel(
         test_config.num_experts,
         test_config.top_k,
         test_config.quantization,
-        test_config.reduce_results,
         test_config.backend,
         _test_body_regular,
         use_shared_experts=test_config.use_shared_experts,
@@ -1578,7 +1520,6 @@ def _parallel_worker(
                 test_config.num_experts,
                 test_config.top_k,
                 test_config.quantization,
-                test_config.reduce_results,
                 test_config.backend,
                 functools.partial(
                     _test_body_config, test_config=test_config, cpu_group=cpu_group
@@ -1597,7 +1538,7 @@ def _parallel_worker(
             failed = failed + 1
             if verbosity > 0:
                 traceback.print_exc()
-                print(f"\n{str(ex)}\nFAILED {ex.__class__}")
+                print(f"\n{str(ex)}\nFAILED")
             else:
                 print("F", end="")
         finally:

tests/kernels/moe/test_ocp_mx_moe.py

@@ -548,7 +548,9 @@ def test_trtllm_gen_mxfp4_fused_moe(
         hidden_states, hidden_states_scale = mxfp8_quantize(
             hidden_states, is_sf_swizzled_layout=False
         )
-        hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(-1)
+        hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
+            *hidden_states.shape[:-1], -1
+        )
     else:
         hidden_states_scale = None
 
@@ -595,20 +597,20 @@ def test_trtllm_gen_mxfp4_fused_moe(
     if beta is not None:
         beta = torch.full((num_experts,), beta, device=hidden_states.device)
     tg_result = tg_mxfp4_moe(
-        router_logits,
-        topk,
-        num_experts,
-        intermediate_size,
-        hidden_size,
-        hidden_states,
-        hidden_states_scale,
-        w13,
-        w13_scale,
-        bias13,
-        w2,
-        w2_scale,
-        bias2,
-        act_type,
+        router_logits=router_logits,
+        topk=topk,
+        num_experts=num_experts,
+        intermediate_size=intermediate_size,
+        hidden_size=hidden_size,
+        hidden_states=hidden_states,
+        hidden_states_scale=hidden_states_scale,
+        w13_weight=w13,
+        w13_weight_scale=w13_scale,
+        w13_bias=bias13,
+        w2_weight=w2,
+        w2_weight_scale=w2_scale,
+        w2_bias=bias2,
+        act_type=act_type,
         alpha=alpha,
         beta=beta,
         limit=limit,

tests/kernels/moe/test_shared_fused_moe_routed_transform.py

@@ -165,7 +165,6 @@ def test_routed_input_transform_inside_vs_outside(
             top_k=top_k,
             hidden_size=latent_size,
             intermediate_size=intermediate_size,
-            reduce_results=False,
             renormalize=True,
             params_dtype=dtype,
             tp_size=1,
@@ -183,7 +182,6 @@ def test_routed_input_transform_inside_vs_outside(
             top_k=top_k,
             hidden_size=latent_size,
             intermediate_size=intermediate_size,
-            reduce_results=False,
             renormalize=True,
             params_dtype=dtype,
             tp_size=1,
@@ -212,34 +210,20 @@ def test_routed_input_transform_inside_vs_outside(
         hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype)
         router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
 
-        # Clone inputs so any in-place modification by Method A
-        # cannot affect Method B's computation.
-        hidden_states_A = hidden_states.clone()
-        router_logits_A = router_logits.clone()
-
         with set_forward_context(None, vllm_config, num_tokens=num_tokens):
-            shared_out_A, routed_out_A = moe_with_transform(
-                hidden_states_A, router_logits_A
-            )
+            # Method A: combined output (shared + routed)
+            combined_A = moe_with_transform(hidden_states, router_logits)
 
+            # Method B: manually transform, get routed output, add shared
             transformed_hidden = routed_transform(hidden_states)
-            shared_out_B, routed_out_B = moe_without_transform(
-                transformed_hidden, router_logits
-            )
+            routed_out_B = moe_without_transform(transformed_hidden, router_logits)
+            shared_out_B = shared_experts(hidden_states)
+            combined_B = shared_out_B + routed_out_B
 
-        expected_shared_out = shared_experts(hidden_states)
-
-        _assert_close(
-            routed_out_A,
-            routed_out_B,
+        torch.testing.assert_close(
+            combined_A,
+            combined_B,
             atol=1e-3,
             rtol=1e-3,
-            label="Routed output: transform inside vs outside",
-        )
-        _assert_close(
-            shared_out_A,
-            expected_shared_out,
-            atol=1e-3,
-            rtol=1e-3,
-            label="Shared expert output",
+            msg="Combined output should match: transform inside vs outside",
         )

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py

@@ -6,6 +6,7 @@ from transformers import AutoModel
 
 from tests.models.utils import check_embeddings_close
 from vllm import TokensPrompt
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
     with vllm_runner(
         model,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         max_model_len=128,
         max_num_batched_tokens=chunk_size,
         enforce_eager=True,

tests/models/language/pooling/test_bge_m3.py

@@ -3,7 +3,6 @@
 import httpx
 import openai
 import pytest
-import pytest_asyncio
 import torch
 
 from ....utils import RemoteOpenAIServer
@@ -25,29 +24,42 @@ sentences_2 = [
 similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
 lexical_score_reference = [0.19554901123046875, 0.0]
 colbert_score_reference = [0.7797, 0.4620]
+SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
+
+
+@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
+def pooling_task(request):
+    yield request.param
 
 
 @pytest.fixture(scope="module")
-def server():
+def server(pooling_task):
     args = [
         "--max-model-len",
         str(MAX_MODEL_LEN),
         "--hf-overrides",
         '{"architectures": ["BgeM3EmbeddingModel"]}',
+        "--pooler-config.task",
+        pooling_task,
     ]
 
     with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
 
 
-@pytest_asyncio.fixture
-async def client(server):
-    async with server.get_async_client() as async_client:
-        yield async_client
-
-
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "embed":
+        with pytest.raises(openai.InternalServerError):
+            await run_client_embeddings(
+                client,
+                MODEL_NAME,
+                sentences_1,
+            )
+        return
+
     embeddings_list_1 = await run_client_embeddings(
         client,
         MODEL_NAME,
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_classify":
+        with pytest.raises(openai.BadRequestError):
+            await sparse_embeddings(client, sentences_1)
+        return
+
     embeddings_1 = await sparse_embeddings(client, sentences_1)
     embeddings_2 = await sparse_embeddings(client, sentences_2)
 
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding_corner_case(
-    client: openai.AsyncOpenAI,
-):
+async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
+    if pooling_task != "token_classify":
+        return
+
+    client = server.get_async_client()
     embeddings = await sparse_embeddings(client, ["Hi"])
     assert len(embeddings) == 1
     assert 2673 in embeddings[0]
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_multi_vector(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_embed":
+        with pytest.raises(openai.BadRequestError):
+            await client.post(
+                "../pooling",
+                body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
+                cast_to=httpx.Response,
+            )
+        return
+
     result_1 = await client.post(
         "../pooling",
         body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},

tests/models/language/pooling/test_extract_hidden_states.py

@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from vllm import TokensPrompt
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
         max_model_len=128,
         enforce_eager=True,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         enable_prefix_caching=True,
     ) as vllm_model:
         pooling_outputs = vllm_model.llm.encode(
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
             assert len(output.prompt_token_ids) == n
             assert len(output.outputs.data) == n
             assert output.num_cached_tokens == 0
-
-        # skip_reading_prefix_cache can still write to cache
-        # to accelerate following requests
-        pooling_outputs = vllm_model.llm.encode(
-            [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
-            pooling_task="embed",
-        )
-
-        for n, output in zip(n_prompt_tokens, pooling_outputs):
-            assert len(output.prompt_token_ids) == n
-            assert output.num_cached_tokens > 0

tests/models/language/pooling/test_multi_vector_retrieval.py

@@ -5,6 +5,7 @@ import torch
 from transformers import AutoModel
 
 from tests.models.utils import check_embeddings_close
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
     with vllm_runner(
         model,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         max_model_len=None,
     ) as vllm_model:
         vllm_outputs = vllm_model.token_embed(example_prompts)

tests/models/language/pooling/test_pooler_config_init_behaviour.py

@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
         model,
         max_model_len=512,
         dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=False),
+        pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
     ) as vllm_model:
         wo_normalize = vllm_model.token_embed(example_prompts)
 
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
         model,
         max_model_len=512,
         dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=True),
+        pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
     ) as vllm_model:
         w_normalize = vllm_model.token_embed(example_prompts)
 

tests/models/multimodal/generation/test_common.py

@@ -86,6 +86,29 @@ COMMON_BROADCAST_SETTINGS = {
 # which cases would be selected and deselected by pytest. In general,
 # this is a good idea for checking your command first, since tests are slow.
 
+
+def _granite4_vision_vllm_to_hf_output(vllm_output, model):
+    """Post-processor for granite4_vision vLLM output.
+
+    Self-contained to avoid calling AutoConfig/AutoTokenizer without
+    trust_remote_code (needed while the model is not in upstream HF).
+    """
+    output_ids, output_str, out_logprobs = vllm_output
+    mm_token_id = 100352
+    hf_output_ids = [
+        token_id
+        for idx, token_id in enumerate(output_ids)
+        if token_id != mm_token_id or idx == 0 or output_ids[idx - 1] != mm_token_id
+    ]
+    hf_output_str = (
+        output_str[1:] if output_str and output_str[0] == " " else output_str
+    )
+    eos_token_id = 100257
+    if hf_output_ids and hf_output_ids[-1] == eos_token_id:
+        hf_output_str = hf_output_str + "<|end_of_text|>"
+    return hf_output_ids, hf_output_str, out_logprobs
+
+
 VLM_TEST_SETTINGS = {
     #### Core tests to always run in the CI
     "llava": VLMTestInfo(
@@ -492,6 +515,20 @@ VLM_TEST_SETTINGS = {
         auto_cls=AutoModelForImageTextToText,
         marks=[large_gpu_mark(min_gb=32)],
     ),
+    "granite4_vision": VLMTestInfo(
+        models=["ibm-granite/granite-vision-4.1-4b"],
+        test_type=(VLMTestType.IMAGE),
+        prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n",
+        max_model_len=8192,
+        auto_cls=AutoModelForImageTextToText,
+        vllm_output_post_proc=_granite4_vision_vllm_to_hf_output,
+        image_size_factors=[(1.0,)],
+        vllm_runner_kwargs={
+            "enable_lora": True,
+            "max_lora_rank": 256,
+            "default_mm_loras": {"image": "ibm-granite/granite-vision-4.1-4b"},
+        },
+    ),
     "h2ovl": VLMTestInfo(
         models=[
             "h2oai/h2ovl-mississippi-800m",

tests/models/registry.py

@@ -899,6 +899,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "zai-org/GLM-ASR-Nano-2512",
         min_transformers_version="5.0.0",
     ),
+    "Granite4VisionForConditionalGeneration": _HfExamplesInfo(
+        "ibm-granite/granite-vision-4.1-4b",
+        is_available_online=False,
+    ),
     "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"),
     "GraniteSpeechForConditionalGeneration": _HfExamplesInfo(
         "ibm-granite/granite-speech-3.3-2b",

tests/quantization/test_gfx950_moe.py

@@ -2,5 +2,5 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
-def test_mi3xx_moe():
-    print("TODO: add tests for Mi3xx MoE quantization")
+def test_mi355_moe():
+    print("TODO: add tests for Mi355 MoE quantization")

tests/reasoning/test_gptoss_reasoning_parser.py

@@ -280,3 +280,72 @@ class TestGptOssStructuralTags:
             assert tag["content"]["type"] == "any_text"
             assert tag["end"] == "<|end|>"
             assert tag["begin"].startswith("<|channel|>")
+
+
+@pytest.mark.parametrize(
+    "output, is_reasoning_end",
+    [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
+)
+def test_gptoss_is_reasoning_end_streaming(
+    output,
+    is_reasoning_end,
+    gpt_oss_tokenizer,
+):
+    """Streaming override must agree with is_reasoning_end for all cases."""
+    tokens = gpt_oss_tokenizer.tokenize(output)
+    parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
+    output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
+    delta_ids = output_ids[-1:] if output_ids else []
+    actual = parser.is_reasoning_end_streaming(output_ids, delta_ids)
+    assert is_reasoning_end == actual
+
+
+@pytest.mark.parametrize(
+    "output, is_reasoning_end",
+    [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
+)
+def test_gptoss_is_reasoning_end_streaming_long_prefix(
+    output,
+    is_reasoning_end,
+    gpt_oss_tokenizer,
+):
+    """Windowing must produce correct results even with a long prefix."""
+    tokens = gpt_oss_tokenizer.tokenize(output)
+    parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
+    output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
+    # Prepend 10k dummy reasoning tokens to simulate a long generation
+    long_prefix = [1] * 10_000
+    padded_ids = long_prefix + list(output_ids)
+    delta_ids = output_ids[-1:] if output_ids else []
+    actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids)
+    assert is_reasoning_end == actual
+
+
+@pytest.mark.parametrize(
+    "output, is_reasoning_end",
+    [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
+)
+def test_gptoss_is_reasoning_end_streaming_large_delta(
+    output,
+    is_reasoning_end,
+    gpt_oss_tokenizer,
+):
+    """Simulate speculative decoding where the entire test sequence arrives
+    as a single large delta appended after a long prefix.  The window must
+    expand to cover delta_ids so the end pattern is never missed."""
+    tokens = gpt_oss_tokenizer.tokenize(output)
+    parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
+    output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
+    long_prefix = [1] * 10_000
+    padded_ids = long_prefix + list(output_ids)
+    # delta_ids = the entire test sequence (as if accepted in one spec step)
+    delta_ids = list(output_ids)
+    actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids)
+    assert is_reasoning_end == actual
+
+
+def test_gptoss_is_reasoning_end_streaming_signature(gpt_oss_tokenizer):
+    """Verify the method is callable with the expected signature."""
+    parser = GptOssReasoningParser(gpt_oss_tokenizer)
+    result = parser.is_reasoning_end_streaming([], [])
+    assert result is False

tests/renderers/inputs/test_preprocess.py

@@ -1,6 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from vllm.renderers.inputs.preprocess import prompt_to_seq
+import pytest
+
+from vllm.renderers.inputs.preprocess import (
+    parse_dec_only_prompt,
+    parse_enc_dec_prompt,
+    prompt_to_seq,
+)
 
 
 def test_empty_input():
@@ -39,3 +45,23 @@ def test_dict_input():
         {"prompt": "foo"},
         {"prompt_token_ids": [1, 2]},
     ]
+
+
+def test_parse_dec_only_prompt_rejects_non_string_prompt_field():
+    with pytest.raises(TypeError, match="Prompt text should be a string"):
+        parse_dec_only_prompt({"prompt": [1, 2, 3], "cache_salt": "abc"})
+
+
+def test_parse_dec_only_prompt_rejects_non_string_prompt_list():
+    with pytest.raises(TypeError, match="Prompt text should be a string"):
+        parse_dec_only_prompt({"prompt": [1, "x"]})
+
+
+def test_parse_enc_dec_prompt_rejects_nested_non_string_prompt_field():
+    with pytest.raises(TypeError, match="Prompt text should be a string"):
+        parse_enc_dec_prompt(
+            {
+                "encoder_prompt": {"prompt": [1, 2, 3]},
+                "decoder_prompt": {"prompt": [4, 5]},
+            }
+        )

tests/transformers_utils/test_hf_overrides_model_type.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test that hf_overrides model_type returns the correct config class."""
+
+import json
+import tempfile
+
+from transformers import PretrainedConfig
+
+from vllm.transformers_utils.config import _CONFIG_REGISTRY, get_config
+
+
+class _TestCustomConfig(PretrainedConfig):
+    model_type = "test_custom_model"
+
+    def __init__(self, custom_attr=42, **kw):
+        super().__init__(**kw)
+        self.custom_attr = custom_attr
+
+
+def test_hf_overrides_model_type_returns_correct_config_class():
+    """When hf_overrides sets model_type to a registered custom type whose
+    checkpoint has a *different* model_type on disk, get_config() must return
+    an instance of the registered config class — not the class that matches
+    the on-disk model_type."""
+
+    # Register the custom config
+    _CONFIG_REGISTRY["test_custom_model"] = _TestCustomConfig
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Checkpoint says model_type="mixtral" on disk
+            cfg = {
+                "model_type": "mixtral",
+                "hidden_size": 128,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 4,
+                "intermediate_size": 256,
+                "num_local_experts": 4,
+                "num_experts_per_tok": 2,
+            }
+            with open(f"{tmpdir}/config.json", "w") as f:
+                json.dump(cfg, f)
+
+            config = get_config(
+                tmpdir,
+                trust_remote_code=False,
+                hf_overrides_kw={
+                    "model_type": "test_custom_model",
+                },
+            )
+
+            from transformers import AutoConfig
+            from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+            # get_config() returns the registered custom class
+            assert isinstance(config, _TestCustomConfig), (
+                f"Expected _TestCustomConfig, got {type(config).__name__}"
+            )
+
+            # AutoConfig has _TestCustomConfig registered under both
+            # the overridden model_type and the on-disk model_type
+            assert CONFIG_MAPPING["test_custom_model"] is _TestCustomConfig
+            assert CONFIG_MAPPING["mixtral"] is _TestCustomConfig
+
+            # AutoConfig.from_pretrained now returns _TestCustomConfig
+            # for this checkpoint (even though its on-disk model_type
+            # is "mixtral")
+            auto_config = AutoConfig.from_pretrained(tmpdir)
+            assert isinstance(auto_config, _TestCustomConfig), (
+                f"Expected _TestCustomConfig from AutoConfig, got "
+                f"{type(auto_config).__name__}"
+            )
+    finally:
+        _CONFIG_REGISTRY.pop("test_custom_model", None)
+        # Restore the original mixtral AutoConfig mapping to avoid
+        # side effects on other tests in the same process
+        from transformers import AutoConfig, MixtralConfig
+
+        AutoConfig.register("mixtral", MixtralConfig, exist_ok=True)

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -10,8 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
 from vllm.distributed.kv_events import BlockRemoved, BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
+    OffloadingConnectorScheduler,
+)
 from vllm.v1.core.kv_cache_utils import BlockHash
-from vllm.v1.kv_offload.abstract import OffloadingEvent
+from vllm.v1.kv_offload.abstract import (
+    OffloadingEvent,
+    OffloadingManager,
+    ReqContext,
+    get_offload_block_hash,
+)
 from vllm.v1.request import RequestStatus
 
 
@@ -105,8 +114,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
         lambda keys, req_context: generate_store_output([])
     )
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
-    runner.manager.lookup.assert_called()
-    assert len(list(runner.manager.lookup.call_args.args[0])) == 1
+    runner.manager.lookup.assert_called_once()
 
     # single block lookup with a hit
     runner.scheduler.reset_prefix_cache()
@@ -114,7 +122,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = (
         lambda keys, req_context: generate_store_output([])
     )
-    runner.manager.lookup.return_value = 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
     )
@@ -126,7 +134,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = (
         lambda keys, req_context: generate_store_output([])
     )
-    runner.manager.lookup.return_value = 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
     )
@@ -210,7 +218,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
 
     # request should now return from preemption
     # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
-    runner.manager.lookup.return_value = 3
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
     runner.manager.prepare_store.side_effect = (
         lambda keys, req_context: generate_store_output(keys)
     )
@@ -251,7 +259,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -263,7 +271,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
 
     # start a new request to load the same first block
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -311,7 +319,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -336,3 +344,137 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
 
     # assert request is deleted
     assert req_id not in runner.scheduler.requests
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for _maximal_prefix_lookup / _sliding_window_lookup
+# ---------------------------------------------------------------------------
+
+
+def _make_scheduler_with_lookup(
+    lookup_results: dict[int, bool | None],
+) -> OffloadingConnectorScheduler:
+    """Create an OffloadingConnectorScheduler with a mocked manager.lookup."""
+    manager = MagicMock(spec=OffloadingManager)
+    manager.lookup.side_effect = lambda key, req_context: lookup_results.get(
+        int(get_offload_block_hash(key).decode()), False
+    )
+
+    scheduler = object.__new__(OffloadingConnectorScheduler)
+    scheduler.manager = manager
+    return scheduler
+
+
+_EMPTY_REQ_CTX = ReqContext()
+
+
+class TestMaximalPrefixLookup:
+    def test_all_hit(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: True})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
+
+    def test_all_miss(self):
+        sched = _make_scheduler_with_lookup({})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+
+    def test_partial_prefix(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: True})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+
+    def test_miss_then_hit(self):
+        sched = _make_scheduler_with_lookup({2: True})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+
+    def test_single_hit(self):
+        sched = _make_scheduler_with_lookup({1: True})
+        assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
+
+    def test_empty(self):
+        sched = _make_scheduler_with_lookup({})
+        assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
+
+    def test_none_defers(self):
+        sched = _make_scheduler_with_lookup({1: None, 2: True})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+
+    def test_none_after_hit_defers(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: None})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+
+    def test_none_stops_at_miss(self):
+        """None is treated as hit for iteration, but miss stops the scan."""
+        sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True})
+        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        # lookup should have been called for blocks 1 and 2 (stops at miss)
+        assert sched.manager.lookup.call_count == 2
+
+
+class TestSlidingWindowLookup:
+    def test_all_hit_exact_window(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: True})
+        assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2
+
+    def test_all_miss(self):
+        sched = _make_scheduler_with_lookup({})
+        assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0
+
+    def test_window_at_end(self):
+        sched = _make_scheduler_with_lookup({2: True, 3: True})
+        assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3
+
+    def test_window_in_middle(self):
+        sched = _make_scheduler_with_lookup({2: True, 3: True})
+        assert (
+            sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3
+        )
+
+    def test_no_full_window_falls_back_to_prefix(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: True})
+        assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2
+
+    def test_single_block_window(self):
+        sched = _make_scheduler_with_lookup({2: True, 3: True})
+        assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3
+
+    def test_gap_resets_consecutive(self):
+        sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True})
+        # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3]
+        assert (
+            sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX)
+            == 3
+        )
+
+    def test_window_prefers_rightmost(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True})
+        # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4
+        # scans right-to-left, finds [4,5] first
+        assert (
+            sched._sliding_window_lookup(to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX)
+            == 5
+        )
+
+    def test_prefix_fallback_with_gap(self):
+        sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True})
+        # window of 4 not found contiguously (gap at 1)
+        assert (
+            sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX)
+            == 1
+        )
+
+    def test_empty(self):
+        sched = _make_scheduler_with_lookup({})
+        assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0
+
+    def test_none_defers(self):
+        sched = _make_scheduler_with_lookup({1: True, 2: None})
+        assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None
+
+    def test_none_with_full_window_still_defers(self):
+        """Even if a real window is found after a None, result is deferred."""
+        # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window
+        # but block 3 was None so defer_lookup is set
+        sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True})
+        assert (
+            sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)
+            is None
+        )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -56,8 +56,12 @@ from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 
 
-def to_keys(int_ids: list[int]) -> list[OffloadKey]:
-    return [make_offload_key(str(i).encode(), 0) for i in int_ids]
+def to_key(int_hash: int) -> OffloadKey:
+    return make_offload_key(str(int_hash).encode(), 0)
+
+
+def to_keys(int_hashes: list[int]) -> list[OffloadKey]:
+    return [to_key(i) for i in int_hashes]
 
 
 class MockLoadStoreSpec(LoadStoreSpec):
@@ -116,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec):
         self.manager = MagicMock(spec=OffloadingManager)
         self.manager.lookup.return_value = 0
         self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
+        self.manager.lookup.return_value = False
         self.handler = MockOffloadingHandler()
 
     def get_manager(self) -> OffloadingManager:
@@ -228,14 +233,14 @@ class RequestRunner:
         self.scheduler_connector: OffloadingConnector = scheduler_connector
 
         # extract mocked OffloadingManager of scheduler connector
-        connector_scheduler = scheduler_connector.connector_scheduler
-        assert connector_scheduler is not None
-        manager = connector_scheduler.manager
+        self.connector_scheduler = scheduler_connector.connector_scheduler
+        assert self.connector_scheduler is not None
+        manager = self.connector_scheduler.manager
         assert isinstance(manager, MagicMock)
         self.manager: MagicMock = manager
 
-        assert len(connector_scheduler.config.kv_group_configs) == 1
-        kv_group_config = connector_scheduler.config.kv_group_configs[0]
+        assert len(self.connector_scheduler.config.kv_group_configs) == 1
+        kv_group_config = self.connector_scheduler.config.kv_group_configs[0]
         assert kv_group_config.gpu_block_size == gpu_block_size
         assert kv_group_config.offloaded_block_size == offloaded_block_size
 

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -631,7 +631,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
         mock_thread.return_value.is_alive.return_value = False
 
         worker.use_mla = True
-        worker.kv_topo.is_mla = True
+        worker.transfer_topo.is_mla = True
 
         # MLA cache tensor: shape[-2] is the block size.
         mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
@@ -692,9 +692,9 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
         # Override TP rank/size to simulate P TP=2
         prefill_worker.tp_rank = P_TP_RANK
         prefill_worker.tp_size = P_TP_SIZE
-        # Update shared dict so kv_topo sees correct TP size
         prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
-        prefill_worker.kv_topo.tp_rank = P_TP_RANK
+        prefill_worker.transfer_topo.tp_rank = P_TP_RANK
+        prefill_worker.transfer_topo.tp_size = P_TP_SIZE
 
         prefill_worker.kv_caches_base_addr = [0x1000]
         prefill_worker.block_len_per_layer = [local_block_len]
@@ -714,7 +714,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
         send_meta.ready.set()
 
         # Compute target D ranks using the production code path
-        target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)
+        target_d_ranks = prefill_worker.transfer_topo.handshake_target_ranks(d_tp_size)
 
         mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
         mock_socket.send_multipart = AsyncMock()

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -21,7 +21,7 @@ from vllm import LLM
 from vllm.config import KVTransferConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     KVOutputAggregator,
-    TpKVTopology,
+    TransferTopology,
     get_current_attn_backend,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
@@ -463,19 +463,20 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         test_shape = self.attn_backends[0].get_kv_cache_shape(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
-        self.kv_topo = TpKVTopology(
+        self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
+            tp_size=self.world_size,
+            block_size=self.block_size,
             engine_id=self.engine_id,
-            remote_tp_size=self._tp_size,  # shared state
-            remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
+            is_mamba=False,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
             attn_backends=self.attn_backends,
             tensor_shape=test_shape,
         )
 
         self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
+            self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
         )
 
     def _nixl_handshake(
@@ -496,7 +497,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         # Adjust remote block length metadata to satisfy heterogeneous TP
         # invariants enforced during handshake validation.
         remote_block_lens = list(self.block_len_per_layer)
-        tp_ratio = self.kv_topo.tp_ratio(remote_tp_size)
+        tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
         if remote_tp_size > self.world_size:
             # P TP > D TP case, block_len of remote is smaller
             remote_block_lens = [
@@ -731,8 +732,9 @@ class TestNixlHandshake:
             assert set(remote_agents.keys()) == set(range(tp_ratio))
 
             remote_engine_id = worker.REMOTE_ENGINE_ID
-            assert worker._tp_size[remote_engine_id] == remote_tp_size
-            assert -tp_ratio == worker.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
+            remote_info = worker.transfer_topo.get_engine_info(remote_engine_id)
+            assert remote_info.remote_tp_size == remote_tp_size
+            assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
             # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
             assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio
             assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio
@@ -796,7 +798,7 @@ class TestNixlHandshake:
             (conn_p0.connector_worker, conn_p1.connector_worker)
         ):
             worker.world_size = p_tp_size
-            worker.kv_topo.remote_tp_size = {worker.engine_id: p_tp_size}
+            worker.transfer_topo.tp_size = p_tp_size
             worker.tp_rank = rank
             worker.use_mla = True
 
@@ -2337,7 +2339,7 @@ def test_compatibility_hash_validation(
         remote_hash = compute_nixl_compatibility_hash(
             remote_vllm_config,
             decode_worker.backend_name,
-            decode_worker.kv_topo.cross_layers_blocks,
+            decode_worker.transfer_topo.cross_layers_blocks,
         )
 
     prefill_block_size = config_overrides.get("block_size", 16)
@@ -2424,12 +2426,13 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     test_shape = backend.get_kv_cache_shape(
         num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
     )
-    decode_worker.kv_topo = TpKVTopology(
+    decode_worker.transfer_topo = TransferTopology(
         tp_rank=decode_worker.tp_rank,
+        tp_size=decode_worker.world_size,
+        block_size=decode_worker.block_size,
         engine_id=decode_worker.engine_id,
-        remote_tp_size=decode_worker._tp_size,  # shared state
-        remote_block_size=decode_worker._block_size,  # shared state
         is_mla=decode_worker.use_mla,
+        is_mamba=False,
         total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
         attn_backends=[backend],
         tensor_shape=test_shape,
@@ -2438,7 +2441,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker.compat_hash = compute_nixl_compatibility_hash(
         decode_worker.vllm_config,
         decode_worker.backend_name,
-        decode_worker.kv_topo.cross_layers_blocks,
+        decode_worker.transfer_topo.cross_layers_blocks,
     )
 
     if error_scenario == "handshake_decode_error":

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -152,13 +152,14 @@ def test_read_blocks_for_req_expands_remote_ids(
 
     remote_engine_id = "remote-engine"
     if has_mamba:
-        worker._mamba_phys_ratio = {remote_engine_id: remote_ratio}
+        worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio}
 
-    # Mock kv_topo: empty remote ranks skips the transfer machinery entirely,
-    # isolating the block-ID expansion logic.
-    worker.kv_topo = MagicMock()
-    worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = []
-    worker.kv_topo.tp_ratio_from_engine_id.return_value = 1
+    # Mock transfer_topo: empty remote ranks skips the transfer machinery
+    # entirely, isolating the block-ID expansion logic.
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.target_remote_ranks.return_value = []
+    worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1)
+    worker.transfer_topo.tp_ratio.return_value = 1
 
     metadata = NixlConnectorMetadata()
     metadata.add_new_req_to_recv(
@@ -317,7 +318,7 @@ def test_get_block_descs_ids_hybrid_ssm():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = 1
-    worker._mamba_phys_ratio = {engine_id: 1}
+    worker._physical_blocks_per_logical = {engine_id: 1}
     worker.block_len_per_layer = [100]
     # num_descs = num_regions * num_blocks (no blocks_first doubling)
     worker.num_descs = 2 * num_blocks
@@ -355,7 +356,7 @@ def test_get_block_descs_ids_kernel_block_mismatch():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = ratio
-    worker._mamba_phys_ratio = {engine_id: ratio}
+    worker._physical_blocks_per_logical = {engine_id: ratio}
     worker.block_len_per_layer = [100]
     worker.num_descs = 2 * num_blocks  # 800
 
@@ -532,15 +533,15 @@ def test_has_mamba_init(
         ((9216, 524288), 4096, 131),
     ],
 )
-def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio):
-    """Verify that compute_mamba_phys_ratio is TP-dependent.
+def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_ratio):
+    """Verify that compute_physical_blocks_per_logical is TP-dependent.
 
     With dimension-sharded Mamba state, the ratio differs across TP sizes
     (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why
-    _mamba_phys_ratio must be stored per-engine.
+    _physical_blocks_per_logical must be stored per-engine.
     """
     from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
-        compute_mamba_phys_ratio,
+        compute_physical_blocks_per_logical,
     )
 
-    assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio
+    assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio

tests/v1/kv_offload/test_cpu_manager.py

@@ -35,8 +35,12 @@ class ExpectedPrepareStoreOutput:
     evicted_keys: list[int]
 
 
-def to_keys(int_ids: list[int]) -> list[OffloadKey]:
-    return [make_offload_key(str(i).encode(), 0) for i in int_ids]
+def to_key(int_hash: int) -> OffloadKey:
+    return make_offload_key(str(int_hash).encode(), 0)
+
+
+def to_keys(int_hashes: list[int]) -> list[OffloadKey]:
+    return [to_key(i) for i in int_hashes]
 
 
 def verify_store_output(
@@ -136,7 +140,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
     manager.complete_store(to_keys([2, 3, 4, 5]))
 
     # block 2 must still be present in the cache
-    assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
+    assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
 
 
 def test_cpu_manager():
@@ -160,7 +164,8 @@ def test_cpu_manager():
     )
 
     # lookup [1, 2] -> not ready
-    assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+    assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
+    assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False
 
     # no events so far
     assert list(cpu_manager.take_events()) == []
@@ -170,9 +175,9 @@ def test_cpu_manager():
     verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
 
     # lookup [1, 2]
-    assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
-    assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
-    assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+    assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False
 
     # prepare store [2, 3, 4, 5] -> evicts [1]
     prepare_store_output = cpu_manager.prepare_store(
@@ -196,6 +201,14 @@ def test_cpu_manager():
     # complete store [2, 3, 4, 5]
     cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
 
+    # lookup (now that we have [2, 3, 4, 5])
+    assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
+    assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False
+
     # prepare load [2, 3]
     prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
     verify_load_output(prepare_load_output, [1, 2])
@@ -238,8 +251,8 @@ def test_cpu_manager():
     cpu_manager.complete_store(to_keys([7, 9]), success=False)
 
     # assert [7] is still stored, but [9] is not
-    assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1
-    assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0
+    assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True
+    assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False
 
     verify_events(
         cpu_manager.take_events(),
@@ -284,7 +297,8 @@ class TestARCPolicy:
         )
 
         # lookup [1, 2] -> not ready
-        assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
+        assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False
 
         # no events so far
         assert list(cpu_manager.take_events()) == []
@@ -294,9 +308,9 @@ class TestARCPolicy:
         verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
 
         # lookup [1, 2]
-        assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
-        assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
-        assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+        assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True
+        assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
+        assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False
 
         # blocks should be in T1 (recent)
         assert len(arc_policy.t1) == 2
@@ -500,7 +514,7 @@ class TestARCPolicy:
         cpu_manager.complete_store(to_keys([5]), success=False)
 
         # block 5 should not be in cache
-        assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0
+        assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False
         # block 5 should not be in T1 or T2
         assert to_keys([5])[0] not in arc_policy.t1
         assert to_keys([5])[0] not in arc_policy.t2
@@ -541,8 +555,8 @@ class TestARCPolicy:
         cpu_manager.complete_store(to_keys([6]))
 
         # verify blocks 2, 3 (in T2) are still present
-        assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
-        assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1
+        assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
+        assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True
 
         # verify events
         events = list(cpu_manager.take_events())
@@ -562,7 +576,8 @@ def test_filter_reused_manager():
     )
 
     # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
-    assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
+    assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False
 
     # prepare store [1, 2] -> should be filtered
     prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -570,7 +585,7 @@ def test_filter_reused_manager():
     assert prepare_store_output.keys_to_store == []
 
     # Lookup [1] -> 2nd time, eligible now
-    assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
 
     # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
     prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -579,12 +594,13 @@ def test_filter_reused_manager():
 
     # Lookup [3, 4] -> 1st time
     # (evicts [2] from tracker since max_size is 3 and tracker has [1])
-    assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False
+    assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False
     # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
     assert to_keys([2])[0] not in manager.counts
 
     # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
-    assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False
     # Verify [2] was re-added with count=1 (not eligible yet)
     assert manager.counts.get(to_keys([2])[0]) == 1
 

tools/pre_commit/generate_attention_backend_docs.py

@@ -235,11 +235,10 @@ def _resolve_import_to_file(
 
 
 def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
-    """Find a compute capability from is_device_capability*() calls in a function.
+    """Find a compute capability from is_device_capability_family() calls in a function.
 
-    Handles two patterns:
-    - is_device_capability_family(N): "M.x" (e.g. 100 -> "10.x")
-    - is_device_capability(N): "M.m" (e.g. 100 -> "10.0")
+    Looks for the pattern: current_platform.is_device_capability_family(N)
+    and converts N (e.g. 100) to a CC string (e.g. "10.x").
     """
     for node in ast.walk(tree):
         if not isinstance(node, ast.FunctionDef) or node.name != func_name:
@@ -248,15 +247,12 @@ def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
             if (
                 isinstance(n, ast.Call)
                 and isinstance(n.func, ast.Attribute)
+                and n.func.attr == "is_device_capability_family"
                 and n.args
                 and isinstance(n.args[0], ast.Constant)
                 and isinstance(n.args[0].value, int)
             ):
-                val = n.args[0].value
-                if n.func.attr == "is_device_capability_family":
-                    return f"{val // 10}.x"
-                elif n.func.attr == "is_device_capability":
-                    return f"{val // 10}.{val % 10}"
+                return f"{n.args[0].value // 10}.x"
     return None
 
 

vllm/_aiter_ops.py

@@ -376,6 +376,22 @@ def _rocm_aiter_fused_topk_fake(
 
 # Cache whether aiter supports FP8 MLA parameters
 _AITER_MLA_SUPPORTS_FP8: bool | None = None
+_AITER_HAS_FUSED_QK_RMSNORM: bool | None = None
+
+
+def check_aiter_fused_qk_rmsnorm() -> bool:
+    """Check if aiter provides fused_qk_rmsnorm (requires AITer >= PR #2442)."""
+    global _AITER_HAS_FUSED_QK_RMSNORM
+    if _AITER_HAS_FUSED_QK_RMSNORM is None:
+        try:
+            from aiter.ops.fused_qk_norm_rope_cache_quant import (  # noqa: F401
+                fused_qk_rmsnorm,
+            )
+
+            _AITER_HAS_FUSED_QK_RMSNORM = True
+        except (ImportError, ModuleNotFoundError, AttributeError):
+            _AITER_HAS_FUSED_QK_RMSNORM = False
+    return _AITER_HAS_FUSED_QK_RMSNORM
 
 
 def _check_aiter_mla_fp8_support() -> bool:
@@ -962,6 +978,44 @@ def _rocm_aiter_triton_add_rmsnorm_pad_fake(
     return out, residual_out
 
 
+def _fused_mla_dual_rms_norm_impl(
+    x1: torch.Tensor,
+    x1_weight: torch.Tensor,
+    x2: torch.Tensor,
+    x2_weight: torch.Tensor,
+    x1_epsilon: float,
+    x2_epsilon: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    try:
+        from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
+    except (ImportError, ModuleNotFoundError) as exc:
+        raise ImportError(
+            "fused_qk_rmsnorm requires a newer AITer version "
+            "(>= PR #2442). Please upgrade aiter or disable the "
+            "fuse_mla_dual_rms_norm pass."
+        ) from exc
+
+    return fused_qk_rmsnorm(
+        q=x1,
+        q_weight=x1_weight,
+        q_eps=x1_epsilon,
+        k=x2,
+        k_weight=x2_weight,
+        k_eps=x2_epsilon,
+    )
+
+
+def _fused_mla_dual_rms_norm_fake(
+    x1: torch.Tensor,
+    x1_weight: torch.Tensor,
+    x2: torch.Tensor,
+    x2_weight: torch.Tensor,
+    x1_epsilon: float,
+    x2_epsilon: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return (torch.empty_like(x1), torch.empty_like(x2))
+
+
 def _rocm_aiter_gemm_a8wfp4_impl(
     x: torch.Tensor,
     w: torch.Tensor,
@@ -1491,6 +1545,13 @@ class rocm_aiter_ops:
                 fake_impl=_triton_rotary_embedding_fake,
             )
 
+            direct_register_custom_op(
+                op_name="fused_mla_dual_rms_norm",
+                op_func=_fused_mla_dual_rms_norm_impl,
+                mutates_args=[],
+                fake_impl=_fused_mla_dual_rms_norm_fake,
+            )
+
             _OPS_REGISTERED = True
 
     @staticmethod
@@ -1537,6 +1598,10 @@ class rocm_aiter_ops:
     def get_triton_rotary_embedding_op() -> OpOverload:
         return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default
 
+    @staticmethod
+    def get_fused_mla_dual_rms_norm_op() -> OpOverload:
+        return torch.ops.vllm.fused_mla_dual_rms_norm.default
+
     @staticmethod
     def rms_norm(
         x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float

vllm/benchmarks/datasets/datasets.py

@@ -1373,26 +1373,6 @@ class ShareGPTDataset(BenchmarkDataset):
         return samples
 
 
-class _ValidateDatasetArgs(argparse.Action):
-    """Argparse action to validate dataset name and path compatibility."""
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        setattr(namespace, self.dest, values)
-
-        # Get current values of both dataset_name and dataset_path
-        dataset_name = getattr(namespace, "dataset_name", "random")
-        dataset_path = getattr(namespace, "dataset_path", None)
-
-        # Validate the combination
-        if dataset_name == "random" and dataset_path is not None:
-            parser.error(
-                "Cannot use 'random' dataset with --dataset-path. "
-                "Please specify the appropriate --dataset-name (e.g., "
-                "'sharegpt', 'custom', 'sonnet') for your dataset file: "
-                f"{dataset_path}"
-            )
-
-
 def add_dataset_parser(parser: FlexibleArgumentParser):
     parser.add_argument(
         "--trust-remote-code",
@@ -1410,7 +1390,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
         "--dataset-name",
         type=str,
         default="random",
-        action=_ValidateDatasetArgs,
         choices=[
             "sharegpt",
             "burstgpt",
@@ -1436,7 +1415,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
         "--dataset-path",
         type=str,
         default=None,
-        action=_ValidateDatasetArgs,
         help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
         "using HF dataset.",
     )
@@ -1608,7 +1586,9 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
         "repetition dataset.",
     )
 
-    speed_bench_group = parser.add_argument_group("speed bench dataset options")
+    speed_bench_group = parser.add_argument_group(
+        "speed bench dataset options", description=SpeedBench.__doc__
+    )
     speed_bench_group.add_argument(
         "--speed-bench-dataset-subset",
         type=str,
@@ -3603,11 +3583,15 @@ class MMStarDataset(HuggingFaceDataset):
 
 class SpeedBench(CustomDataset):
     """
-    Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
+    SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
+
     Download the dataset using:
-    curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -
+
+    `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -`
     """  # noqa: E501
 
+    DOWNLOAD_SCRIPT_URL = "https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py"
+
     def __init__(self, **kwargs) -> None:
         self.dataset_subset = kwargs.pop("dataset_subset", "qualitative")
         self.category = kwargs.pop("category", None)
@@ -3618,6 +3602,13 @@ class SpeedBench(CustomDataset):
         if self.dataset_path is None:
             raise ValueError("dataset_path must be provided for loading data.")
 
+        if not Path(self.dataset_path).is_dir():
+            raise ValueError(
+                f"dataset_path {self.dataset_path} is not a directory. "
+                f"Please make sure to download the dataset from HuggingFace using "
+                f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
+            )
+
         self.data = []
 
         # Load the JSONL file
@@ -3628,7 +3619,11 @@ class SpeedBench(CustomDataset):
 
         # check if the JSONL file has a 'turns' column
         if "messages" not in jsonl_data.columns:
-            raise ValueError("JSONL file must contain a 'messages' column.")
+            raise ValueError(
+                "JSONL file must contain a 'messages' column. "
+                "Please make sure to download the dataset from HuggingFace using "
+                f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
+            )
 
         for _, row in jsonl_data.iterrows():
             # sample only from a specific category if specified

vllm/benchmarks/serve.py

@@ -1711,12 +1711,25 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
             trust_remote_code=args.trust_remote_code,
         )
 
+    # Validate dataset name/path
     if args.dataset_name is None:
         raise ValueError(
             "Please specify '--dataset-name' and the corresponding "
             "'--dataset-path' if required."
         )
 
+    if (
+        args.dataset_name
+        in ["random", "random-mm", "random-rerank", "prefix_repetition"]
+        and args.dataset_path is not None
+    ):
+        raise ValueError(
+            f"Cannot use '{args.dataset_name}' dataset with --dataset-path. "
+            "Please specify the appropriate --dataset-name (e.g., "
+            "'sharegpt', 'custom', 'sonnet') for your dataset file: "
+            f"{args.dataset_path}"
+        )
+
     # Map general --input-len and --output-len to all dataset-specific arguments
     if args.input_len is not None:
         args.random_input_len = args.input_len

vllm/compilation/compiler_interface.py

@@ -295,6 +295,9 @@ class InductorStandaloneAdaptor(CompilerInterface):
             },
         }
 
+        if is_torch_equal_or_newer("2.13.0.dev"):
+            compile_kwargs["donate_graph_module"] = True  # type: ignore[assignment]
+
         use_aot: bool = supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT
         # only add 'aot' parameter if both supported and enabled...
         # this will set bundled_autograd_cache

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -7,6 +8,7 @@ import torch._inductor.pattern_matcher as pm
 from torch import fx
 from torch._inductor.pattern_matcher import PatternMatcherPass
 
+import vllm.ir.ops
 import vllm.model_executor.layers.quantization.utils.fp8_utils  # noqa: F401
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
@@ -20,7 +22,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.platforms import current_platform
 
 from ..inductor_pass import enable_fake_mode
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
+from ..vllm_inductor_pass import (
+    VllmFusionPatternMatcherPass,
+    VllmInductorPass,
+    VllmPatternMatcherPass,
+    VllmPatternReplacement,
+)
 from .act_quant_fusion import ActivationQuantPattern
 from .matcher_utils import (
     MatcherFusedAddRMSNorm,
@@ -512,3 +519,101 @@ class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass):
 
     def uuid(self) -> str:
         return VllmInductorPass.hash_source(self, AddAiterRMSNormPadPattern)
+
+
+class MLADualRMSNormPattern(
+    VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
+):
+    """
+    Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into
+    AITER's ``fused_qk_rmsnorm`` HIP kernel.
+
+    Target FX-graph pattern (unfused, ``vllm_ir`` stage)::
+
+        gemm -> split_with_sizes([q_dim, kv_dim])
+            +-- q_c     -> vllm_ir.rms_norm(q_c, q_w, eps)
+            +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim])
+                            +-- kv_c -> vllm_ir.rms_norm(kv_c, kv_w, eps)
+                            +-- k_pe
+
+    The pattern covers the connected subgraph rooted at the first
+    ``split_with_sizes`` (which produces ``q_c`` and ``kv_lora``),
+    through the two ``rms_norm`` calls, and the ``k_pe`` passthrough.
+    """
+
+    def __init__(self, epsilon: float) -> None:
+        self._epsilon = epsilon
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        q_dim, kv_c_dim, k_pe_dim = 8, 4, 2
+        return [
+            self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim),
+            self.empty_bf16(q_dim),
+            self.empty_bf16(kv_c_dim),
+        ]
+
+    @property
+    def pattern(
+        self,
+    ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        eps = self._epsilon
+
+        def _pattern(
+            projected: torch.Tensor,
+            q_weight: torch.Tensor,
+            kv_weight: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            q_dim = q_weight.shape[0]
+            kv_dim = projected.shape[-1] - q_dim
+            kv_c_dim = kv_weight.shape[0]
+            k_pe_dim = kv_dim - kv_c_dim
+            q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)
+            kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)
+            q_normed = vllm.ir.ops.rms_norm(q_c, q_weight, eps)
+            kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps)
+            return q_normed, kv_normed, k_pe
+
+        return _pattern
+
+    @property
+    def replacement(
+        self,
+    ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        eps = self._epsilon
+
+        def _replacement(
+            projected: torch.Tensor,
+            q_weight: torch.Tensor,
+            kv_weight: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            q_dim = q_weight.shape[0]
+            kv_dim = projected.shape[-1] - q_dim
+            kv_c_dim = kv_weight.shape[0]
+            k_pe_dim = kv_dim - kv_c_dim
+            q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)
+            kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)
+            q_normed, kv_normed = torch.ops.vllm.fused_mla_dual_rms_norm(
+                q_c,
+                q_weight,
+                kv_c,
+                kv_weight,
+                eps,
+                eps,
+            )
+            return q_normed, kv_normed, k_pe
+
+        return _replacement
+
+
+class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass):
+    """
+    Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in
+    MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's
+    ``fused_qk_rmsnorm`` HIP kernel.
+    """
+
+    def __init__(self, config: VllmConfig) -> None:
+        super().__init__(config, "mla_dual_rms_norm_fusion_pass")
+
+        for epsilon in [1e-5, 1e-6]:
+            self.register(MLADualRMSNormPattern(epsilon))

vllm/compilation/passes/inductor_pass.py

@@ -51,6 +51,15 @@ def pass_context(compile_range: Range) -> Generator[None, None, None]:
         _pass_context = prev_context
 
 
+@functools.cache
+def _hash_source_cached(*srcs: str | type | types.FunctionType) -> str:
+    hasher = hashlib.sha256()
+    for src in srcs:
+        src_str = src if isinstance(src, str) else inspect.getsource(src)
+        hasher.update(src_str.encode("utf-8"))
+    return hasher.hexdigest()
+
+
 class InductorPass(CustomGraphPass):  # type: ignore[misc]
     """
     A custom graph pass that uses a hash of its source as the UUID.
@@ -72,19 +81,16 @@ class InductorPass(CustomGraphPass):  # type: ignore[misc]
         Utility method to hash the sources of functions or objects.
         :param srcs: strings or objects to add to the hash.
         Objects and functions have their source inspected.
+        Results are cached by resolved types to avoid repeated
+        inspect.getsource() calls.
         :return:
         """
-        hasher = hashlib.sha256()
-        for src in srcs:
-            if isinstance(src, str):
-                src_str = src
-            elif isinstance(src, (types.FunctionType, type)):
-                src_str = inspect.getsource(src)
-            else:
-                # object instance
-                src_str = inspect.getsource(src.__class__)
-            hasher.update(src_str.encode("utf-8"))
-        return hasher.hexdigest()
+        # Resolve instances to their class for a hashable cache key.
+        cache_key = tuple(
+            src if isinstance(src, (str, type, types.FunctionType)) else src.__class__
+            for src in srcs
+        )
+        return _hash_source_cached(*cache_key)
 
     @staticmethod
     def hash_dict(dict_: dict[Any, Any]) -> str:

vllm/compilation/passes/pass_manager.py

@@ -19,6 +19,7 @@ from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
 
 if rocm_aiter_ops.is_enabled():
     from .fusion.rocm_aiter_fusion import (
+        MLADualRMSNormFusionPass,
         RocmAiterRMSNormQuantFusionPass,
         RocmAiterSiluMulFp8GroupQuantFusionPass,
         RocmAiterTritonAddRMSNormPadFusionPass,
@@ -155,6 +156,9 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
             if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
                 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
 
+            if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
+                self.passes += [MLADualRMSNormFusionPass(config)]
+
             if self.pass_config.fuse_rope_kvcache:
                 self.passes += [SplitCoalescingPass(config)]
                 self.passes += [ScatterSplitReplacementPass(config)]

vllm/compilation/piecewise_backend.py

@@ -11,6 +11,7 @@ from typing import Any
 
 import torch._functorch.config
 import torch.fx as fx
+from torch._dynamo.utils import dynamo_timed
 from torch._inductor.runtime.triton_heuristics import CachingAutotuner
 from torch._logging._internal import trace_structured
 
@@ -275,6 +276,7 @@ class PiecewiseBackend:
 
             range_entry.compiled = True
 
+    @dynamo_timed("vllm_log_compile_start_torch_trace_only")
     def _log_compile_start(self, compile_range: Range):
         """Log compilation event for TORCH_TRACE/tlparse."""
         is_cudagraph_size = (

vllm/config/compilation.py

@@ -142,6 +142,8 @@ class PassConfig:
     # ROCm/AITER specific fusions
     fuse_act_padding: bool = None  # type: ignore[assignment]
     """Fuse the custom RMSNorm + padding ops."""
+    fuse_mla_dual_rms_norm: bool = None  # type: ignore[assignment]
+    """Fuse paired q/kv RMS norms in MLA attention."""
     fuse_rope_kvcache: bool = None  # type: ignore[assignment]
     """Fuse the QK rope + KV cache ops."""
 
@@ -224,6 +226,7 @@ class PassConfig:
         "fuse_gemm_comms",
         "fuse_allreduce_rms",
         "fuse_act_padding",
+        "fuse_mla_dual_rms_norm",
         "fuse_rope_kvcache",
         mode="wrap",
     )
@@ -270,6 +273,12 @@ class PassConfig:
                 "The fusion will be disabled."
             )
             self.fuse_act_padding = False
+        if self.fuse_mla_dual_rms_norm and not current_platform.is_rocm():
+            logger.warning_once(
+                "MLA dual RMS norm fusion requires ROCm/AITER. "
+                "The fusion will be disabled."
+            )
+            self.fuse_mla_dual_rms_norm = False
         if self.fuse_rope_kvcache and not current_platform.is_rocm():
             logger.warning_once(
                 "KV cache fusion currently only enabled on ROCm. "
@@ -526,13 +535,15 @@ class CompilationConfig:
     model's budget range. User-provided positive value overrides
     auto-inference."""
 
-    encoder_cudagraph_max_frames_per_batch: int = 0
+    encoder_cudagraph_max_frames_per_batch: int | None = None
     """Maximum total video frames per batch for encoder CUDA graph capture.
     Controls the cu_seqlens buffer size (one entry per attention sequence,
-    i.e. one per video frame). If 0 (default), auto-inferred per budget
-    level as token_budget (tight bound: packing guarantees
-    sum(T_i) <= token_budget). Positive value overrides auto-inference
-    and applies to all budget levels."""
+    i.e. one per video frame).
+    If None (default), auto-inferred as encoder_cudagraph_max_vision_items_per_batch
+    * max_frames_per_video (model-specific value according to processing_info).
+    Positive value overrides auto-inference and applies to all budget levels.
+    If we limit the video count per prompt to `0`, it will also be set to `0`
+    (i.e., fall back to image-only mode)."""
 
     # Inductor capture
     compile_sizes: list[int | str] | None = None
@@ -984,11 +995,12 @@ class CompilationConfig:
             )
         if (
             self.cudagraph_mm_encoder
+            and self.encoder_cudagraph_max_frames_per_batch is not None
             and self.encoder_cudagraph_max_frames_per_batch < 0
         ):
             raise ValueError(
                 "encoder_cudagraph_max_frames_per_batch must be "
-                "non-negative (0 = auto-infer)"
+                "non-negative (None = auto-infer)"
             )
 
         if self.backend == "":

vllm/config/parallel.py

@@ -36,7 +36,7 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
 DataParallelBackend = Literal["ray", "mp"]
 EPLBPolicyOption = Literal["default"]
 DCPCommBackend = Literal["ag_rs", "a2a"]
-EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "pynccl"]
+EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"]
 All2AllBackend = Literal[
     "naive",
     "pplx",
@@ -90,6 +90,7 @@ class EPLBConfig:
     Backend for EPLB expert weight communication:
     - "torch_nccl": Use torch.distributed on the device process group
     - "torch_gloo": Use torch.distributed gloo with CPU staging
+    - "nixl": Use NIXL/ RIXL with staged send/recv buffers
     - "pynccl": Use PyNccl send/recv
     - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync)
     """