Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/image_build/image_build.sh

@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
 }
 
 ecr_login() {
-    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
-    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
+    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
 }
 
 prepare_cache_tags() {

.buildkite/image_build/image_build_cpu.sh

@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
 
 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then

.buildkite/image_build/image_build_cpu_arm64.sh

@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
 
 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then

.buildkite/intel_jobs/engine_intel.yaml

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

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

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

@@ -1,55 +0,0 @@
-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,6 +240,7 @@ 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}}"
@@ -282,7 +283,6 @@ 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

@@ -1,55 +0,0 @@
-#!/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,15 +29,6 @@ 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/kernels.yaml

@@ -141,7 +141,6 @@ steps:
     - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
     - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
-    - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py

.buildkite/test_areas/spec_decode.yaml

@@ -12,17 +12,6 @@ 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
@@ -34,18 +23,6 @@ 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
@@ -66,17 +43,6 @@ 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

CMakeLists.txt

@@ -952,9 +952,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
@@ -1242,6 +1240,7 @@ endif()
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
     include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/cutlass_fa3.cmake)
     include(cmake/external_projects/qutlass.cmake)
 
     # vllm-flash-attn should be last as it overwrites some CMake functions

benchmarks/kernels/ir/bench_ir_ops.py

@@ -1,378 +0,0 @@
-#!/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

@@ -1,29 +0,0 @@
-# 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,49 +161,16 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
-    # 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()
+    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)
     else()
         message(STATUS "compile riscv with scalar")
-        set(MARCH_FLAGS -march=rv64gc)
+        list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
     endif()
     list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
 else()

cmake/external_projects/cutlass_fa3.cmake

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
+#
+# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
+# as a self-contained extension (_cutlass_fa3_C). This provides a high-
+# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
+#
+# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
+
+# Guard: CUDA >= 12.4 required for SM90a features used by FA3
+if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
+  message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
+  # Create empty target so setup.py doesn't fail on unsupported systems
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+# Guard: SM90 architecture required
+set(CUTLASS_FA3_SUPPORT_ARCHS)
+list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
+cuda_archs_loose_intersection(
+  CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(NOT CUTLASS_FA3_ARCHS)
+  message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+include(FetchContent)
+
+# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
+# We only need the source files, not the build system, so we use
+# FetchContent_Populate to download without building.
+if (DEFINED ENV{SGL_ATTN_SRC_DIR})
+  set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
+endif()
+if(SGL_ATTN_SRC_DIR)
+  FetchContent_Declare(cutlass_fa3
+    SOURCE_DIR ${SGL_ATTN_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_fa3
+    GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
+    GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_fa3)
+if(NOT cutlass_fa3_POPULATED)
+  FetchContent_Populate(cutlass_fa3)
+endif()
+message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
+
+# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
+# to avoid version conflicts). Use FetchContent_Populate to avoid running
+# CUTLASS's own CMakeLists.txt which would create conflicting targets.
+if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+  set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+endif()
+if(CUTLASS_FA3_CUTLASS_SRC_DIR)
+  FetchContent_Declare(cutlass_for_fa3
+    SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_for_fa3
+    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+    GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_for_fa3)
+if(NOT cutlass_for_fa3_POPULATED)
+  FetchContent_Populate(cutlass_for_fa3)
+endif()
+message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
+
+set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
+
+# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
+# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
+# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
+file(GLOB FA3_INSTANTIATION_SOURCES
+  # BF16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
+  # FP16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
+  # FP8 (e4m3) instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
+
+set(FA3_CORE_SOURCES
+  "${FA3_SRC}/flash_api.cpp"
+  "${FA3_SRC}/flash_prepare_scheduler.cu"
+  "${FA3_SRC}/flash_fwd_combine.cu")
+
+set(FA3_ALL_SOURCES
+  "${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
+  ${FA3_CORE_SOURCES}
+  ${FA3_INSTANTIATION_SOURCES})
+
+set(FA3_INCLUDE_DIRS
+  ${FA3_SRC}
+  ${cutlass_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
+  ${CMAKE_CURRENT_SOURCE_DIR}/csrc)
+
+# Set SM90a gencode flags for all FA3 CUDA sources
+set_gencode_flags_for_srcs(
+  SRCS "${FA3_ALL_SOURCES}"
+  CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
+
+define_extension_target(_cutlass_fa3_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${FA3_ALL_SOURCES}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
+  USE_SABI 3
+  WITH_SOABI)
+
+# FA3-specific compile options for CUDA and C++ source files:
+# - C++17 required by CUTLASS
+# - Fast math for performance
+# - Relaxed constexpr for CUTLASS template metaprogramming
+# - Disable backward pass, dropout, uneven K (not needed for inference)
+# - Enable varlen-only mode (all our use cases are variable-length)
+target_compile_options(_cutlass_fa3_C PRIVATE
+  $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
+  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
+
+target_compile_definitions(_cutlass_fa3_C PRIVATE
+  CUTE_USE_PACKED_TUPLE=1
+  CUTLASS_ENABLE_GDC_FOR_SM90
+  CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
+  CUTLASS_ENABLE_TENSOR_CORE_MMA=1
+  FLASHATTENTION_DISABLE_BACKWARD
+  FLASHATTENTION_DISABLE_DROPOUT
+  FLASHATTENTION_DISABLE_UNEVEN_K
+  FLASHATTENTION_DISABLE_SM8x
+  FLASHATTENTION_VARLEN_ONLY)
+
+message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")

csrc/cpu/cpu_types_riscv.hpp

@@ -1,25 +1,832 @@
 #ifndef CPU_TYPES_RISCV_HPP
 #define CPU_TYPES_RISCV_HPP
 
-// 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).
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <iostream>
+#include <limits>
+#include <riscv_vector.h>
+#include <torch/all.h>
 
-#ifndef __riscv_vector
-  #error "cpu_types_riscv.hpp included in a non-RVV translation unit"
+// ============================================================================
+// 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)));
 #endif
 
-#ifndef __riscv_v_min_vlen
-  #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
+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
 
-#include "cpu_types_riscv_defs.hpp"
-#include "cpu_types_riscv_impl.hpp"
+#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
-#endif  // CPU_TYPES_RISCV_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

csrc/cpu/cpu_types_riscv_defs.hpp

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

@@ -1,891 +0,0 @@
-#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/cutlass_fa3_extension.cc

@@ -0,0 +1,72 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
+ * Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
+ *
+ * This registers the FA3 forward pass as a PyTorch C++ extension under
+ * the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
+ *
+ * Original source:
+ *   https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+ *   sgl-kernel/csrc/flash_extension.cc
+ */
+#include <Python.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "sgl_flash_kernel_ops.h"
+
+TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
+  /*
+   * CUTLASS FA3 MLA forward pass.
+   * Signature matches sgl-attn's mha_fwd() exactly.
+   */
+  m.def(
+      "fwd(Tensor   q,"
+      "    Tensor   k,"
+      "    Tensor   v,"
+      "    Tensor?  k_new,"
+      "    Tensor?  v_new,"
+      "    Tensor?  q_v,"
+      "    Tensor?  out,"
+      "    Tensor?  cu_seqlens_q,"
+      "    Tensor?  cu_seqlens_k,"
+      "    Tensor?  cu_seqlens_k_new,"
+      "    Tensor?  seqused_q,"
+      "    Tensor?  seqused_k,"
+      "    int?     max_seqlen_q,"
+      "    int?     max_seqlen_k,"
+      "    Tensor?  page_table,"
+      "    Tensor?  kv_batch_idx,"
+      "    Tensor?  leftpad_k,"
+      "    Tensor?  rotary_cos,"
+      "    Tensor?  rotary_sin,"
+      "    Tensor?  seqlens_rotary,"
+      "    Tensor?  q_descale,"
+      "    Tensor?  k_descale,"
+      "    Tensor?  v_descale,"
+      "    float?   softmax_scale,"
+      "    bool     is_causal,"
+      "    int      window_size_left,"
+      "    int      window_size_right,"
+      "    int      attention_chunk,"
+      "    float    softcap,"
+      "    bool     is_rotary_interleaved,"
+      "    Tensor?  scheduler_metadata,"
+      "    int      num_splits,"
+      "    bool?    pack_gqa,"
+      "    int      sm_margin,"
+      "    Tensor?  sinks"
+      ") -> (Tensor, Tensor, Tensor, Tensor)");
+
+  m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
+}
+
+// Python module initialization for _cutlass_fa3_C
+PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
+  static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
+                                      nullptr, 0, nullptr};
+  return PyModule_Create(&module);
+}

csrc/libtorch_stable/ops.h

@@ -134,13 +134,4 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor& input,
                               torch::stable::Tensor& input_global_scale);
 
-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets);
-
 #endif

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -1,468 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
- * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
- */
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include <cutlass/arch/arch.h>
-
-#include "cutlass_extensions/common.hpp"
-
-#include "cute/tensor.hpp"
-#include "cutlass/tensor_ref.h"
-#include "cutlass/epilogue/collective/default_epilogue.hpp"
-#include "cutlass/epilogue/thread/linear_combination.h"
-#include "cutlass/gemm/dispatch_policy.hpp"
-#include "cutlass/gemm/group_array_problem_shape.hpp"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-
-#include "cutlass/util/packed_stride.hpp"
-#include <cassert>
-
-using namespace cute;
-
-// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
-template <typename ElementAB, typename ElementC, typename ElementSF,
-          typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-__global__ void __mxfp4_get_group_gemm_starts(
-    ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
-    ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
-    LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
-    ElementAB* a_base_as_int, ElementAB* b_base_as_int,
-    ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
-    ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
-    const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
-    int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
-    const int64_t a_stride_val, const int64_t b_stride_val,
-    const int64_t c_stride_val, const int K, const int N) {
-  int64_t expert_id = threadIdx.x;
-  if (expert_id >= gridDim.x * blockDim.x) {
-    return;
-  }
-  int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
-  int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
-  int64_t group_size = 32;
-  int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
-  int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
-  int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
-  assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
-         "unexpected problem sizes");
-
-  int64_t half_k = static_cast<int64_t>(k / 2);
-  int64_t group_k = static_cast<int64_t>(k / group_size);
-  // Shape of A as uint8/byte = [M, K // 2]
-  a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
-  // Shape of B as uint8/byte = [E, N, K // 2]
-  b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
-  // Shape of C = [M, N]
-  out_offsets[expert_id] = out_base_as_int + expert_offset * n;
-  // Shape of a_scale = [sum(sf_sizes), K // group_size]
-  a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
-
-  assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Shape of B scale = [E, N, K // group_size]
-  b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
-  assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Initialize strides
-  a_strides[expert_id] = a_stride_val;
-  b_strides[expert_id] = b_stride_val;
-  c_strides[expert_id] = c_stride_val;
-
-  LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
-  LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
-
-  *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-  *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-}
-
-#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE,               \
-                                       TENSOR_C_TYPE, C_TYPE, LayoutSFA,       \
-                                       LayoutSFB, ScaleConfig)                 \
-  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                       \
-    __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
-                                  LayoutSFB, ScaleConfig>                      \
-        <<<1, num_experts, 0, stream>>>(                                       \
-            static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()),               \
-            static_cast<C_TYPE**>(out_starts.data_ptr()),                      \
-            static_cast<SF_TYPE**>(a_scales_starts.data_ptr()),                \
-            static_cast<SF_TYPE**>(b_scales_starts.data_ptr()),                \
-            reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),               \
-            reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()),               \
-            static_cast<C_TYPE*>(out_tensors.data_ptr()),                      \
-            static_cast<SF_TYPE*>(a_scales.data_ptr()),                        \
-            static_cast<SF_TYPE*>(b_scales.data_ptr()),                        \
-            static_cast<int32_t*>(expert_offsets.data_ptr()),                  \
-            static_cast<int32_t*>(sf_offsets.data_ptr()),                      \
-            static_cast<int32_t*>(problem_sizes.data_ptr()),                   \
-            static_cast<int64_t*>(a_strides.data_ptr()),                       \
-            static_cast<int64_t*>(b_strides.data_ptr()),                       \
-            static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val,         \
-            b_stride_val, c_stride_val, K, N);                                 \
-  }
-
-template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-void mxfp4_run_get_group_gemm_starts(
-    const torch::stable::Tensor& a_starts,
-    const torch::stable::Tensor& b_starts,
-    const torch::stable::Tensor& out_starts,
-    const torch::stable::Tensor& a_scales_starts,
-    const torch::stable::Tensor& b_scales_starts,
-    const torch::stable::Tensor& layout_sfa,
-    const torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a_strides,
-    const torch::stable::Tensor& b_strides,
-    const torch::stable::Tensor& c_strides, int64_t a_stride_val,
-    int64_t b_stride_val, int64_t c_stride_val,
-    torch::stable::Tensor const& a_tensors,
-    torch::stable::Tensor const& b_tensors,
-    torch::stable::Tensor const& out_tensors,
-    torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales,
-    torch::stable::Tensor const& expert_offsets,
-    torch::stable::Tensor const& sf_offsets,
-    torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
-  int num_experts = (int)expert_offsets.size(0);
-  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
-
-  STD_TORCH_CHECK(out_tensors.size(1) == N,
-                  "Output tensor shape doesn't match expected shape");
-  STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
-                  "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
-                  " dimension must match");
-  if (false) {
-  }
-  // MXFP4 uses E8M0 (float_ue8m0_t) scale factors
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::BFloat16,
-                                 cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
-                                 ScaleConfig)
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::Half, half,
-                                 LayoutSFA, LayoutSFB, ScaleConfig)
-  else {
-    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
-  }
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm_sm100(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  using ProblemShape =
-      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
-  using ElementType = cutlass::float_e2m1_t;
-  using ElementSFType = cutlass::float_ue8m0_t;
-  using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-  using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-
-  using ElementC = OutType;
-  using ElementD = ElementC;
-  using ElementAccumulator = float;
-  // Layout definitions
-  using LayoutA = cutlass::layout::RowMajor;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = LayoutC;
-
-  static constexpr int AlignmentA = 32;
-  static constexpr int AlignmentB = 32;
-  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-
-  // Architecture definitions
-  using ArchTag = cutlass::arch::Sm100;
-  using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
-  using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-  using StageCountType = cutlass::gemm::collective::StageCountAuto;
-
-  using ClusterShape = Shape<_1, _1, _1>;
-  struct MMA1SMConfig {
-    using MmaTileShape = Shape<_128, _128, _128>;
-    using KernelSchedule =
-        cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
-    using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
-  };
-
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
-          ClusterShape, Shape<_128, _64>, ElementAccumulator,
-          ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
-          LayoutC*, AlignmentD,
-          typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
-
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
-          ElementB, LayoutB*, AlignmentB, ElementAccumulator,
-          typename MMA1SMConfig::MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
-
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using ScaleConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-
-  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-
-  torch::stable::Tensor a_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor out_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor a_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor layout_sfa = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor layout_sfb = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor a_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor c_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-
-  mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
-      a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
-      layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
-      b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
-      b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
-
-  // Create an instance of the GEMM
-  Gemm gemm_op;
-
-  UnderlyingProblemShape* problem_sizes_as_shapes =
-      static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  // Set the Scheduler info
-  cutlass::KernelHardwareInfo hw_info;
-  using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
-      PersistentTileSchedulerSm100GroupParams<
-          typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
-  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
-  scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device_index();
-  static std::unordered_map<int, int> cached_sm_counts;
-  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
-    cached_sm_counts[hw_info.device_id] =
-        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
-            hw_info.device_id);
-  }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
-
-  // Mainloop Arguments
-  typename GemmKernel::MainloopArguments mainloop_args{
-      static_cast<const ElementType**>(a_ptrs.data_ptr()),
-      static_cast<StrideA*>(a_strides1.data_ptr()),
-      static_cast<const ElementType**>(b_ptrs.data_ptr()),
-      static_cast<StrideB*>(b_strides1.data_ptr()),
-      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
-
-  // Epilogue Arguments
-  typename GemmKernel::EpilogueArguments epilogue_args{
-      {},  // epilogue.thread
-      nullptr,
-      static_cast<StrideC*>(c_strides1.data_ptr()),
-      static_cast<ElementD**>(out_ptrs.data_ptr()),
-      static_cast<StrideC*>(c_strides1.data_ptr())};
-  auto& fusion_args = epilogue_args.thread;
-  // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
-  fusion_args.alpha_ptr = nullptr;
-  fusion_args.beta_ptr = nullptr;
-  fusion_args.alpha = 1.0f;
-  fusion_args.alpha_ptr_array = nullptr;
-  fusion_args.dAlpha = {_0{}, _0{}, 0};
-  fusion_args.beta = 0.0f;
-  fusion_args.beta_ptr_array = nullptr;
-  fusion_args.dBeta = {_0{}, _0{}, 0};
-
-  // Gemm Arguments
-  typename GemmKernel::Arguments args{
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, problem_sizes_as_shapes, nullptr},
-      mainloop_args,
-      epilogue_args,
-      hw_info,
-      scheduler};
-
-  size_t workspace_size = Gemm::get_workspace_size(args);
-  auto workspace =
-      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
-                           std::nullopt, a.device());
-  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
-
-  auto can_implement_status = gemm_op.can_implement(args);
-  STD_TORCH_CHECK(
-      can_implement_status == cutlass::Status::kSuccess,
-      "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
-
-  // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize MXFP4 GEMM: status=", (int)status,
-                  " workspace_size=", workspace_size,
-                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
-
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to run MXFP4 GEMM");
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  int32_t version_num = get_sm_version_num();
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  if (version_num >= 100 && version_num < 120) {
-    run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-    return;
-  }
-#endif
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
-      version_num, ". Required capability: 100");
-}
-
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
-// E8M0 scale factors stored as uint8
-constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
-#endif
-
-#define CHECK_TYPE(x, st, m)             \
-  STD_TORCH_CHECK(x.scalar_type() == st, \
-                  ": Inconsistency of torch::stable::Tensor type:", m)
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
-#define CHECK_INPUT(x, st, m) \
-  CHECK_TH_CUDA(x, m);        \
-  CHECK_CONTIGUOUS(x, m);     \
-  CHECK_TYPE(x, st, m)
-
-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets) {
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  // Input validation
-  CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
-  CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
-  // MXFP4 uses E8M0 scale factors (stored as uint8)
-  CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
-  CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
-
-  STD_TORCH_CHECK(
-      a_blockscale.dim() == 2,
-      "expected a_blockscale to be of shape [num_experts, rounded_m,"
-      " k // group_size], observed rank: ",
-      a_blockscale.dim())
-  STD_TORCH_CHECK(b_blockscales.dim() == 3,
-                  "expected b_blockscale to be of shape: "
-                  " [num_experts, n, k // group_size], observed rank: ",
-                  b_blockscales.dim())
-  STD_TORCH_CHECK(problem_sizes.dim() == 2,
-                  "problem_sizes must be a 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have the shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32.");
-
-  int M = static_cast<int>(a.size(0));
-  int N = static_cast<int>(b.size(1));
-  int E = static_cast<int>(b.size(0));
-  int K = static_cast<int>(2 * b.size(2));
-
-  if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  } else {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  }
-#else
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
-      "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
-}

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -1,432 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 activation quantization kernel for MoE experts.
- * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
- * over 32-element groups.
- *
- * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
- *   - Block size 32 (2 threads per SF in PACK16 mode)
- *   - E8M0 (power-of-two) scale factors
- *   - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
- */
-
-// MXFP4 requires PACK16 mode (16 elements per thread) so that
-// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
-// Must be defined before any header that (transitively) includes
-// nvfp4_utils.cuh.
-#define NVFP4_ENABLE_ELTS16 1
-
-#include <cuda.h>
-#include <cuda_runtime_api.h>
-#include <cuda_runtime.h>
-#include <cuda_fp8.h>
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-#include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
-#include "cuda_utils.h"
-
-#include "nvfp4_utils.cuh"
-static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
-              "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
-
-#include "launch_bounds_utils.h"
-
-namespace vllm {
-
-// MXFP4 block size constants
-static constexpr int MXFP4_SF_VEC_SIZE = 32;
-
-// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
-// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
-static constexpr int MXFP4_NUM_THREADS_PER_SF =
-    MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
-
-// MXFP4 quantization kernel for experts.
-// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
-// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
-// SiLU(gate)*up before quantization.
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts, bool low_latency) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    if constexpr (SMALL_NUM_EXPERTS) {
-      for (int i = 0; i < n_experts; i++) {
-        uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
-        uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
-        if (rowIdx >= current_offset && rowIdx < next_offset) {
-          rowIdx_in_expert = rowIdx - current_offset;
-          expert_idx = i;
-          break;
-        }
-      }
-    } else {
-      uint32_t local_offsets[17];
-      for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
-        *reinterpret_cast<int4*>(local_offsets) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start]));
-        *reinterpret_cast<int4*>(local_offsets + 4) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 4]));
-        *reinterpret_cast<int4*>(local_offsets + 8) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 8]));
-        *reinterpret_cast<int4*>(local_offsets + 12) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 12]));
-        local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
-
-#pragma unroll
-        for (int i = 0; i < 16; i++) {
-          if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
-            rowIdx_in_expert = rowIdx - local_offsets[i];
-            expert_idx = chunk_start + i;
-            break;
-          }
-        }
-      }
-    }
-
-    // Load input and optionally apply fused SiLU+Mul
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    // Block E8M0 scales only; no extra tensor-level scale in this path
-    constexpr float SFScaleVal = 1.0f;
-    // UE8M0_SF=true for MXFP4 E8M0 scale factors
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-// Large M_topk variant using shared memory for expert offsets
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128)
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  extern __shared__ uint32_t shared_input_offsets[];
-
-  if constexpr (SMALL_NUM_EXPERTS) {
-    for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
-      shared_input_offsets[i] = input_offset_by_experts[i];
-    }
-  } else {
-    for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
-      *reinterpret_cast<int4*>(&shared_input_offsets[i]) =
-          *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
-    }
-    if (threadIdx.x == 0) {
-      shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
-    }
-  }
-
-  __syncthreads();
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    // Binary search through experts using shared memory
-    int left = 0, right = n_experts - 1;
-    while (left <= right) {
-      int mid = (left + right) / 2;
-      uint32_t mid_offset = shared_input_offsets[mid];
-      uint32_t next_offset = shared_input_offsets[mid + 1];
-
-      if (rowIdx >= mid_offset && rowIdx < next_offset) {
-        rowIdx_in_expert = rowIdx - mid_offset;
-        expert_idx = mid;
-        break;
-      } else if (rowIdx < mid_offset) {
-        right = mid - 1;
-      } else {
-        left = mid + 1;
-      }
-    }
-
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    // MXFP4 has no global scale - only block-level E8M0 scale factors
-    constexpr float SFScaleVal = 1.0f;
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-template <typename T, bool FUSE_SILU_MUL = false>
-void mxfp4_quant_impl(void* output, void* output_scale, void* input,
-                      void* input_offset_by_experts,
-                      void* output_scale_offset_by_experts, int m_topk, int k,
-                      int n_experts, cudaStream_t stream) {
-  int multiProcessorCount =
-      get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
-
-  int const workSizePerRow = k / ELTS_PER_THREAD;
-  int const totalWorkSize = m_topk * workSizePerRow;
-  dim3 block(std::min(workSizePerRow, 512));
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
-  dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
-                     multiProcessorCount * numBlocksPerSM));
-  while (grid.x <= multiProcessorCount && block.x > 64) {
-    grid.x *= 2;
-    block.x = (block.x + 1) / 2;
-  }
-
-  int const blockRepeat =
-      (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
-  if (blockRepeat > 1) {
-    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
-    if (n_experts >= 4) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    }
-  } else {
-    if (n_experts >= 16) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, 0, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts, /* bool low_latency */ true);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
-          m_topk, k, reinterpret_cast<T*>(input),
-          reinterpret_cast<fp4_packed_t*>(output),
-          reinterpret_cast<uint32_t*>(output_scale),
-          reinterpret_cast<uint32_t*>(input_offset_by_experts),
-          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-          n_experts, /* bool low_latency */ true);
-    }
-  }
-}
-
-}  // namespace vllm
-
-/*Quantization entry for mxfp4 experts quantization*/
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
-#define CHECK_INPUT(x, m) \
-  CHECK_TH_CUDA(x, m);    \
-  CHECK_CONTIGUOUS(x, m);
-
-constexpr auto HALF = torch::headeronly::ScalarType::Half;
-constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
-constexpr auto INT = torch::headeronly::ScalarType::Int;
-constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
-
-static constexpr int MXFP4_BLOCK_SIZE = 32;
-
-static void validate_mxfp4_experts_quant_inputs(
-    torch::stable::Tensor const& output,
-    torch::stable::Tensor const& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts, int64_t m_topk, int64_t k) {
-  CHECK_INPUT(output, "output");
-  CHECK_INPUT(output_scale, "output_scale");
-  CHECK_INPUT(input, "input");
-  CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
-  CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
-
-  STD_TORCH_CHECK(output.dim() == 2);
-  STD_TORCH_CHECK(output_scale.dim() == 2);
-  STD_TORCH_CHECK(input.dim() == 2);
-  STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
-
-  STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
-  // output is uint8 (two mxfp4 values packed into one uint8)
-  // output_scale is int32 (four E8M0 values packed into one int32)
-  STD_TORCH_CHECK(output.scalar_type() == UINT8);
-  STD_TORCH_CHECK(output_scale.scalar_type() == INT);
-
-  STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
-  STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output.size(0) == m_topk);
-  STD_TORCH_CHECK(output.size(1) == k / 2);
-  int scales_k = k / MXFP4_BLOCK_SIZE;
-  // K-dimension scale columns padded to a multiple of 4 for swizzle layout
-  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
-  // 4 = 4 E8M0 values packed into one int32
-  STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
-}
-
-void mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k = input.size(1);
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-void silu_and_mul_mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k_times_2 = input.size(1);
-  STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
-  auto k = k_times_2 / 2;
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
-// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
-// .cpp files and cannot gate the registration from there.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
-  m.impl("silu_and_mul_mxfp4_experts_quant",
-         TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
-}

csrc/libtorch_stable/torch_bindings.cpp

@@ -116,12 +116,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
       " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
 
-  // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
-  ops.def(
-      "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
-      " Tensor a_blockscale, Tensor b_blockscales,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
-
   // Compute NVFP4 block quantized tensor.
   ops.def(
       "scaled_fp4_quant(Tensor input,"
@@ -155,19 +149,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
       "Tensor output_scale_offset_by_experts) -> ()");
 
-  // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
-  ops.def(
-      "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
-  // Fused SiLU+Mul+MXFP4 experts quantization.
-  ops.def(
-      "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
-      "output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
   // Fused SiLU+Mul+NVFP4 quantization.
   ops.def(
       "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -252,8 +233,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("silu_and_mul_scaled_fp4_experts_quant",
            TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
   ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
-  // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
-  // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
+
+  // W4A8 ops: impl registrations are in the source files
+  // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
 #endif
 }
 

csrc/moe/topk_softmax_kernels.cu

@@ -126,9 +126,7 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        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;
+        const float softmax_val = expf(val - float_max) * normalizing_factor;
         output[idx] = softmax_val;
     }
 }
@@ -149,9 +147,7 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        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;
+        const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
         output[idx] = sigmoid_val;
     }
 }
@@ -446,19 +442,6 @@ __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

csrc/sgl_flash_kernel_ops.h

@@ -0,0 +1,45 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
+ * Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
+ * NO MODIFICATIONS from the original (except removing unused macros).
+ */
+
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Tensor.h>
+#include <torch/library.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "sgl_kernel_torch_shim.h"
+
+/*
+ * From flash-attention (sgl-attn fork)
+ */
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
+    at::Tensor q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
+    at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
+    std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
+    std::optional<at::Tensor> q_v_,  // MLA value projection query
+    std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> cu_seqlens_k_,
+    std::optional<at::Tensor> cu_seqlens_k_new_,
+    std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
+    std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
+    std::optional<at::Tensor> page_table_,
+    std::optional<at::Tensor> kv_batch_idx_,
+    std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
+    std::optional<at::Tensor> rotary_sin_,
+    std::optional<at::Tensor> seqlens_rotary_,
+    std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
+    std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
+    bool is_causal, int64_t window_size_left, int64_t window_size_right,
+    int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
+    std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
+    std::optional<bool> pack_gqa_, int64_t sm_margin,
+    std::optional<const at::Tensor>& sinks_);

csrc/sgl_kernel_torch_shim.h

@@ -0,0 +1,121 @@
+/* Adapted from:
+ * https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
+ * Provides make_pytorch_shim() template for PyTorch op registration type
+ * conversion. NO MODIFICATIONS from the original.
+ */
+
+#pragma once
+
+#include <torch/library.h>
+
+/**
+ * Unfortunately, the type signatures of the flash_attn ops are not compatible
+ * with the PyTorch library bindings. To get around that we use
+ * `make_pytorch_shim` which creates a lambda that exposes the API using
+ * PyTorch compatible types to the types, then converts them to the types
+ * expected by the flash_attn ops. This shims allows us to make minimal changes
+ * to `flash_api.cpp` making it easier to synchronize with upstream changes.
+ *
+ * The `pytorch_library_compatible_type` struct is used to map from the
+ * flash_attn ops types to a PyTorch library compatible one. The main issues is
+ * that the following types are not support by PyTorch library bindings:
+ *  - `int`
+ *  - `float`
+ *  - `std::optional<T> &`
+ *  - `std::optional<const at::Tensor> &`
+ * So we convert them to (respectively):
+ *  - `int64_t`
+ *  - `double`
+ *  - `const std::optional<T>&`
+ *  - `const std::optional<at::Tensor>&`
+ */
+
+template <typename T>
+struct pytorch_library_compatible_type {
+  using type = T;
+  static T convert_from_type(T arg) { return arg; }
+};
+
+template <typename T>
+using pytorch_library_compatible_type_t =
+    typename pytorch_library_compatible_type<T>::type;
+
+template <typename T>
+T convert_from_pytorch_compatible_type(
+    pytorch_library_compatible_type_t<T> arg) {
+  return pytorch_library_compatible_type<T>::convert_from_type(arg);
+}
+
+// Map `c10::optional<T> &` -> `const c10::optional<T>&`
+//  (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
+//   the optional container)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>&> {
+  using type = const c10::optional<T>&;
+  static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
+    return const_cast<c10::optional<T>&>(arg);
+  }
+};
+
+// Map `c10::optional<T>` ->
+//          `c10::optional<pytorch_library_compatible_type_t<T>>`
+//  (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>> {
+  using type = c10::optional<pytorch_library_compatible_type_t<T>>;
+  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
+      c10::optional<T> arg) {
+    return arg;
+  }
+};
+
+// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
+template <>
+struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
+  using type = const c10::optional<at::Tensor>&;
+  static c10::optional<const at::Tensor>& convert_from_type(
+      const c10::optional<at::Tensor>& arg) {
+    return const_cast<c10::optional<const at::Tensor>&>(
+        reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
+  }
+};
+
+// Map `int` -> `int64_t`
+template <>
+struct pytorch_library_compatible_type<int> {
+  using type = int64_t;
+  static int convert_from_type(int64_t arg) {
+    TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
+                "int64_t value is too large to be converted to int");
+    TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
+                "int64_t value is too small to be converted to int");
+    return arg;
+  }
+};
+
+// Map `float` -> `double`
+template <>
+struct pytorch_library_compatible_type<float> {
+  using type = double;
+  static float convert_from_type(double arg) {
+    TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
+                "double value is too large to be converted to float");
+    return arg;
+  }
+};
+
+//
+//  Shim Utils
+//
+
+template <typename Ret, typename... Args>
+auto make_pytorch_shim(Ret (*fun)(Args... args)) {
+  return [fun](pytorch_library_compatible_type_t<Args>... args) {
+    return fun(convert_from_pytorch_compatible_type<Args>(args)...);
+  };
+}

docker/Dockerfile

@@ -580,12 +580,33 @@ 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
-ARG FLASHINFER_VERSION=0.6.8.post1
+# 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
 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 \
-    && flashinfer download-cubin
+    && 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
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES

docker/Dockerfile.nightly_torch

@@ -217,13 +217,16 @@ 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.8.post1
+# 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
 # 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.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/Dockerfile.rocm

@@ -2,11 +2,6 @@
 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
@@ -197,10 +192,9 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
 FROM base AS build_deep
 ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="5d90af8b"
+ARG DEEPEP_BRANCH="e84464ec"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
-ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
 ENV ROCSHMEM_DIR=/opt/rocshmem
 
 RUN git clone ${ROCSHMEM_REPO} \
@@ -208,11 +202,13 @@ RUN git clone ${ROCSHMEM_REPO} \
  && git checkout ${ROCSHMEM_BRANCH} \
  && mkdir -p projects/rocshmem/build \
  && cd projects/rocshmem/build \
- && bash ../scripts/build_configs/all_backends \
-      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-      -DROCM_PATH=/opt/rocm \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
+ && cmake .. \
+    -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+    -DROCM_PATH=/opt/rocm \
+    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
+    -DUSE_EXTERNAL_MPI=OFF \
+ && make -j \
+ && make install
 
 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
@@ -221,47 +217,6 @@ 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
@@ -307,6 +262,30 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
     && echo "Detected vLLM version: ${VLLM_VERSION}" \
     && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
 
+# Fail if git-based package dependencies are found in requirements files
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
+RUN echo "Checking for git-based packages in requirements files..." \
+    && echo "Checking common.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
+         echo "ERROR: Git-based packages found in common.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in common.txt"; \
+       fi \
+    && echo "Checking rocm.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
+         echo "ERROR: Git-based packages found in rocm.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in rocm.txt"; \
+       fi \
+    && echo "All requirements files are clean - no git-based packages found"
+
 # Pin vLLM dependencies to exact versions of custom ROCm wheels
 # This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
 COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
@@ -340,7 +319,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # -----------------------
 # Test vLLM image
-FROM mori_base AS test
+FROM base AS test
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -431,7 +410,7 @@ RUN printf '%s\n' \
 
 # -----------------------
 # Final vLLM image
-FROM mori_base AS final
+FROM base AS final
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -468,8 +447,6 @@ 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
@@ -487,9 +464,7 @@ 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 \
-    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${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="v1.1.0"
+ARG MORI_BRANCH="2d02c6a9"
 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.8.post1"
+      "default": "0.6.7"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/contributing/model/transcription.md

@@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
 
 The API server takes care of basic audio I/O and optional chunking before building prompts:
 
-- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
+- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`.
 - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
 - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
 
@@ -206,8 +206,8 @@ Relevant server logic:
     async def _preprocess_speech_to_text(...):
         language = self.model_cls.validate_language(request.language)
         ...
-        y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate)
-        duration = get_audio_duration(y=y, sr=sr)
+        y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate)
+        duration = librosa.get_duration(y=y, sr=sr)
         do_split_audio = (self.asr_config.allow_audio_chunking
                         and duration > self.asr_config.max_audio_clip_s)
         chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))

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 `vllm.utils.profiling` helpers are deprecated and will be removed in
-    `v0.21`. Please use Python's `cProfile` module directly instead.
+    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.
 
 ### Example usage - decorator
 

docs/design/attention_backends.md

@@ -132,16 +132,6 @@ Priority is **1 = highest** (tried first).
 | 6 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 7 | `FLASHMLA_SPARSE` |
 
-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -169,7 +159,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.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
 | `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 |
@@ -209,13 +199,14 @@ configuration.
 
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
+| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
-| `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
+### Video inference support (experimental)
 
 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.
 
-    Mixed inputs (image+video) per prompt are also supported now.
+    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.
 
 ## Model integration via `SupportsEncoderCudaGraph`
 
@@ -76,7 +76,6 @@ 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.
@@ -97,7 +96,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 `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).
+* `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).
 
 ## Usage guide
 
@@ -143,6 +142,7 @@ 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,6 +150,7 @@ 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}'
 ```
 
@@ -168,6 +169,7 @@ compilation_config = {
 
 model = vllm.LLM(
     model="Qwen/Qwen3-VL-32B",
+    limit_mm_per_prompt='{"image": 0}',
     compilation_config=compilation_config,
 )
 ```

docs/design/fusions.md

@@ -21,7 +21,6 @@ or just on the low or high end.
 | Fusion                                                                         | `PassConfig` flag            | Fused operations                               | Default at                     | E2E Speedup        | Fullgraph | `num_tokens` |
 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms)                  | `fuse_allreduce_rms`         | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20%              | No        | Low          |
-| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm)                       | `fuse_minimax_qk_norm`       | Q/K variance all-reduce → Q/K RMSNorm          | Off by default                 | 2-3%               | No        | Low          |
 | [Attention + Quant](#attention--quantization-fuse_attn_quant)                  | `fuse_attn_quant`            | Attention output → FP8/NVFP4 quant             | Off by default                 | 3-7%               | Yes       | Always       |
 | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant)              | `fuse_attn_quant`            | MLA Attention output → FP8/NVFP4 quant         | Off by default                 | TBD                | Yes       | Always       |
 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache)             | `fuse_rope_kvcache`          | Rotary embedding → KV cache write              | O2 (ROCm/AITER only)           | 2-4%               | No        | Low          |
@@ -31,7 +30,6 @@ 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
 
@@ -42,7 +40,6 @@ The table below lists the quantization schemes supported by each fusion on each
 | Fusion                       | SM100 (Blackwell)                        | SM90 (Hopper)                            | SM89 (Ada)                               | SM80 (Ampere) | ROCm                                     |
 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
-| `fuse_minimax_qk_norm`\*     | FP16/BF16                                | FP16/BF16                                | FP16/BF16                                | FP16/BF16     | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
 | `fuse_attn_quant` (MLA)\*    | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static(untested)\*                   |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
@@ -52,15 +49,11 @@ 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)
 for per-backend details.
 
-\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
-tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
-
 † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
 other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -191,35 +184,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
 
 - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
 
-### MiniMax QK Norm (`fuse_minimax_qk_norm`)
-
-!!! info
-    This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
-    the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
-    and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
-    optimization level.
-
-**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
-per-token Q/K variances before applying RMS normalization to Q and K.
-
-This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
-`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
-`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
-
-Example:
-
-```bash
-vllm serve MiniMaxAI/MiniMax-M2.5 \
-  --tensor-parallel-size 4 \
-  --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
-```
-
-**Code locations.**
-
-- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
-- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
-- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
-
 ### Sequence Parallelism (`enable_sp`)
 
 **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
@@ -383,44 +347,6 @@ 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,7 +56,6 @@ 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`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
+- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
 - **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/multimodal_inputs.md

@@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl
 Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.
 
 ```python
+import librosa
 from vllm import LLM, SamplingParams
 from vllm.multimodal.audio import split_audio
-from vllm.multimodal.media.audio import load_audio
 
 # Load long audio file
-audio, sr = load_audio("long_audio.wav", sr=16000)
+audio, sr = librosa.load("long_audio.wav", sr=16000)
 
 # Split into chunks at low-energy (quiet) regions
 chunks = split_audio(
@@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows:
         base_url=openai_api_base,
     )
 
-    # Any format supported by soundfile/PyAV is supported
+    # Any format supported by librosa is supported
     audio_url = AudioAsset("winning_call").url
     audio_base64 = encode_base64_content_from_url(audio_url)
 

docs/features/quantization/online.md

@@ -17,9 +17,6 @@ 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:
@@ -27,7 +24,6 @@ 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
@@ -36,7 +32,8 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
 | ------ | ------------- | ------------------ | ----- |
 | `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 | |
-| `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 |
+
+Support for additional schemes will be added in future versions of vllm.
 
 ## Advanced Configuration
 

docs/features/reasoning_outputs.md

@@ -283,7 +283,9 @@ curl http://localhost:8000/v1/chat/completions \
     "messages": [
       { "role": "user", "content": "9.11 and 9.8, which is greater?" }
     ],
-    "thinking_token_budget": 10
+    "extra_body": {
+      "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 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.
+`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.
 
 ### Pooling multitask support
 
-Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
+Pooling multitask support is deprecated and will be removed in v0.20. 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/pooling_models/token_embed.md

@@ -9,14 +9,14 @@
 - Online APIs:
     - Pooling API (`/pooling`)
 
-The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
+The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token.
 
 Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).
 
 !!! note
 
     Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not 
-    what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or
+    what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or
     `--pooler-config.task token_embed` online.
 
 ## Typical Use Cases

docs/models/supported_models.md

@@ -560,7 +560,6 @@ 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` | | |
@@ -683,24 +682,6 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 !!! note
     `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
 
-#### Realtime Transcription
-
-Speech models that support streaming transcription via the
-[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
-WebSocket endpoint.
-
-| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
-| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
-| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
-| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
-
-!!! note
-    `VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
-
-    `Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
-    You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
-    when serving.
-
 ## Pooling Models
 
 See [this page](pooling_models/README.md) for more information on how to use pooling models.

docs/serving/expert_parallel_deployment.md

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

docs/serving/openai_compatible_server.md

@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
 - [Translation API](#translations-api) (`/v1/audio/translations`)
     - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
 - [Realtime API](#realtime-api) (`/v1/realtime`)
-    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
+    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
 
 In addition, we have the following custom APIs:
 

examples/offline_inference/vision_language.py

@@ -394,24 +394,18 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"
 
     prompts = [
         (
@@ -431,7 +425,6 @@ 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,
@@ -441,23 +434,18 @@ def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<vision><|image_pad|></vision>"
-    video_placeholder = "<vision><|video_pad|></vision>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|system|>\nYou are a helpful assistant.<|endofturn|>\n"
-            f"<|user|>\n{placeholder}"
+            f"<|user|>\n<vision>{placeholder}</vision>"
             f"{question}<|endofturn|>\n"
             "<|assistant|>\n"
         )
@@ -578,7 +566,6 @@ 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,
@@ -587,19 +574,14 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
 
     prompts = [
         (
@@ -620,7 +602,6 @@ 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,
@@ -629,20 +610,15 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
 
     prompts = [
         (
@@ -663,7 +639,6 @@ 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,
@@ -672,20 +647,15 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
 
     prompts = [
         (
@@ -706,7 +676,6 @@ 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,
@@ -715,19 +684,14 @@ def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
             "size": {"shortest_edge": 12544, "longest_edge": 47040000},
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
 
     prompts = [
         (
@@ -808,12 +772,11 @@ 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=16384 if modality in ("video", "image+video") else 8192,
-        limit_mm_per_prompt=mm_limit,
+        max_model_len=8192 if modality == "image" else 16384,
+        limit_mm_per_prompt={modality: 1},
     )
 
     messages = list()
@@ -865,29 +828,6 @@ 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}")
 
@@ -936,25 +876,19 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         enforce_eager=True,
     )
 
-    image_placeholder = "<IMG_CONTEXT>"
-    video_placeholder = "<video>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<IMG_CONTEXT>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + "\n" + video_placeholder
+        placeholder = "<video>"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -975,26 +909,20 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         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 = image_placeholder
+        placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -1015,23 +943,17 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<image>"
-    video_placeholder = "<video>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<image>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + "\n" + video_placeholder
+        placeholder = "<video>"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     messages = [
@@ -1088,27 +1010,21 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1125,27 +1041,21 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -1349,26 +1259,22 @@ def run_llava_next_video(questions: list[str], modality: str) -> ModelRequestDat
 
 # LLaVA-OneVision
 def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData:
-    image_placeholder = "<image>"
-    video_placeholder = "<video>"
+    if modality == "video":
+        prompts = [
+            f"<|im_start|>user <video>\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
+    elif modality == "image":
+        prompts = [
+            f"<|im_start|>user <image>\n{question}<|im_end|><|im_start|>assistant\n"
+            for question in questions
+        ]
 
-    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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
     return ModelRequestData(
@@ -1401,7 +1307,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", "image+video"]
+    assert modality in ["image", "video"]
     # If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
 
     # 2.0
@@ -1423,13 +1329,12 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
     # NOTE The stop_token_ids are different for various versions of MiniCPM-V
     # 2.0
@@ -1442,22 +1347,17 @@ 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]
 
-    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
+    modality_placeholder = {
+        "image": "(<image>./</image>)",
+        "video": "(<video>./</video>)",
+    }
 
     prompts = [
         tokenizer.apply_chat_template(
             [
                 {
                     "role": "user",
-                    "content": f"{placeholder}\n{question}",
+                    "content": f"{modality_placeholder[modality]}\n{question}",
                 }
             ],
             tokenize=False,
@@ -1566,24 +1466,20 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
         max_num_batched_tokens=36864,
     )
 
-    image_placeholder = "<|image|>"
-    video_placeholder = "<|video|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video|>"
+    else:
+        raise ValueError(f"Unsupported modality for molmo2: {modality}")
 
     prompts = [
         f"{placeholder}<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1667,25 +1563,19 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "[unused19]"
-    video_placeholder = "[unused32]"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "[unused19]"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "[unused32]"
 
     prompts = [
         (
@@ -1733,25 +1623,18 @@ 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=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
-
-    image_placeholder = "<image>"
-    video_placeholder = "<video>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<image>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + "\n" + video_placeholder
+        placeholder = "<video>"
 
     prompts = [
         f"<|im_start|>user\n\n{placeholder}\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1963,7 +1846,6 @@ 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,
@@ -1973,23 +1855,18 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
             "min_pixels": 28 * 28,
             "max_pixels": 1280 * 28 * 28,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2006,7 +1883,6 @@ 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,
@@ -2016,23 +1892,18 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2049,7 +1920,6 @@ 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,
@@ -2059,18 +1929,13 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_bos|><|IMAGE|><|vision_eos|>"
-    video_placeholder = "<|vision_bos|><|VIDEO|><|vision_eos|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|IMAGE|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|VIDEO|>"
 
     default_system = (
         "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
@@ -2081,7 +1946,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{placeholder}"
+            f"<|im_start|>user\n<|vision_bos|>{placeholder}<|vision_eos|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2097,7 +1962,6 @@ 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,
@@ -2107,23 +1971,18 @@ def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2140,7 +1999,6 @@ 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,
@@ -2150,23 +2008,18 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
             "max_pixels": 1280 * 28 * 28,
             "fps": 1,
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2337,7 +2190,6 @@ 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,
@@ -2345,23 +2197,18 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
             "architectures": ["Tarsier2ForConditionalGeneration"],
             "model_type": "tarsier2",
         },
-        limit_mm_per_prompt=mm_limit,
+        limit_mm_per_prompt={modality: 1},
     )
 
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
     if modality == "image":
-        placeholder = image_placeholder
+        placeholder = "<|image_pad|>"
     elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
+        placeholder = "<|video_pad|>"
 
     prompts = [
         (
             "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
+            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
             f"{question}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
@@ -2510,24 +2357,6 @@ 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)
 
@@ -2610,7 +2439,7 @@ def parse_args():
         "--modality",
         type=str,
         default="image",
-        choices=["image", "video", "image+video", "vision_chunk"],
+        choices=["image", "video", "vision_chunk"],
         help="Modality of the input.",
     )
     parser.add_argument(
@@ -2717,42 +2546,23 @@ 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": _mm_data(data, modality),
-            "multi_modal_uuids": _mm_uuid(uuid, modality),
+            "multi_modal_data": {modality: data},
+            "multi_modal_uuids": {modality: uuid},
         }
         inputs_with_empty_media = {
             "prompt": prompts[0],
-            "multi_modal_data": _mm_empty(modality),
-            "multi_modal_uuids": _mm_uuid(uuid, modality),
+            "multi_modal_data": {modality: None},
+            "multi_modal_uuids": {modality: uuid},
         }
     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,
@@ -2762,7 +2572,7 @@ def main(args):
                 modality,
             )
         else:
-            # Use the same image/video for all prompts
+            # Use the same image for all prompts
             inputs = []
             inputs_with_empty_media = []
             for i in range(args.num_prompts):
@@ -2770,15 +2580,15 @@ def main(args):
                 inputs.append(
                     {
                         "prompt": prompts[i % len(prompts)],
-                        "multi_modal_data": _mm_data(data, modality),
-                        "multi_modal_uuids": _mm_uuid(uuid, modality),
+                        "multi_modal_data": {modality: data},
+                        "multi_modal_uuids": {modality: uuid},
                     }
                 )
                 inputs_with_empty_media.append(
                     {
                         "prompt": prompts[i % len(prompts)],
-                        "multi_modal_data": _mm_empty(modality),
-                        "multi_modal_uuids": _mm_uuid(uuid, modality),
+                        "multi_modal_data": {modality: None},
+                        "multi_modal_uuids": {modality: uuid},
                     }
                 )
 

examples/offline_inference/vision_language_multi_image.py

@@ -310,38 +310,6 @@ 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"
 
@@ -1519,7 +1487,6 @@ 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/online_serving/disaggregated_serving/example_mm_serve.py

@@ -1,117 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Disaggregated multimodal serving: render → generate round-trip.
-
-Demonstrates the two-phase disaggregated flow:
-  1. /v1/chat/completions/render  – preprocesses a multimodal chat request
-     into token IDs and serialized tensor features.
-  2. /inference/v1/generate       – runs inference on the preprocessed tokens.
-
-The render response is passed *directly* to generate with only
-``sampling_params`` added, showing that the two endpoints compose with
-zero client-side transformation.
-
-Launch the server first:
-
-    vllm serve Qwen/Qwen3-VL-2B-Instruct \
-        --dtype bfloat16 --max-model-len 4096 --enforce-eager
-
-Then run this script:
-
-    python example_mm_serve.py
-"""
-
-import io
-
-import pybase64 as base64
-import requests
-from PIL import Image
-from transformers import AutoTokenizer
-
-BASE_URL = "http://localhost:8000"
-MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
-
-
-def make_data_url(image: Image.Image) -> str:
-    """Encode a PIL image as a base64 data URL."""
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    b64 = base64.b64encode(buf.getvalue()).decode()
-    return f"data:image/png;base64,{b64}"
-
-
-def main():
-    # -- Step 1: Create a test image (solid red) -------------------------
-    image = Image.new("RGB", (224, 224), color=(255, 0, 0))
-    data_url = make_data_url(image)
-    print("Created 224x224 red test image")
-
-    # -- Step 2: Render (preprocess) -------------------------------------
-    render_payload = {
-        "model": MODEL_NAME,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": data_url}},
-                    {
-                        "type": "text",
-                        "text": "What color is this image? Answer in one word.",
-                    },
-                ],
-            }
-        ],
-    }
-
-    print("\n--- Render ---")
-    render_resp = requests.post(
-        f"{BASE_URL}/v1/chat/completions/render", json=render_payload
-    )
-    render_resp.raise_for_status()
-    render_data = render_resp.json()
-
-    print(f"Response keys: {list(render_data.keys())}")
-    print(f"Number of token_ids: {len(render_data['token_ids'])}")
-
-    features = render_data.get("features")
-    if features and features.get("kwargs_data"):
-        print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
-        for modality, items in features["kwargs_data"].items():
-            print(
-                f"  {modality}: {len(items)} item(s), "
-                f"first item type: {type(items[0])} length: {len(items[0])}"
-                if items
-                else "First item: (empty)"
-            )
-    else:
-        print("WARNING: no kwargs_data in render response")
-
-    # -- Step 3: Generate (inference) ------------------------------------
-    # Pass the render output directly — only add sampling_params.
-    generate_payload = render_data
-    generate_payload["sampling_params"] = {
-        "max_tokens": 20,
-        "temperature": 0.0,
-    }
-
-    print("\n--- Generate ---")
-    gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
-    gen_resp.raise_for_status()
-    gen_data = gen_resp.json()
-
-    # -- Step 4: Decode & print ------------------------------------------
-    output_ids = gen_data["choices"][0]["token_ids"]
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    text = tokenizer.decode(output_ids, skip_special_tokens=True)
-
-    print(f"Output token count: {len(output_ids)}")
-    print(f"Generated text: {text!r}")
-
-    if "red" in text.lower():
-        print("\nModel correctly identified the red image.")
-    else:
-        print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/openai_chat_completion_client_for_multimodal.py

@@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                     {
                         "type": "input_audio",
                         "input_audio": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                             "data": audio_base64,
                             "format": "wav",
                         },
@@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                     {
                         "type": "audio_url",
                         "audio_url": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                             "url": audio_url
                         },
                     },
@@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                     {
                         "type": "audio_url",
                         "audio_url": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                             "url": f"data:audio/ogg;base64,{audio_base64}"
                         },
                     },

examples/online_serving/openai_realtime_client.py

@@ -12,6 +12,7 @@ model, for example:
 Requirements:
 - vllm with audio support
 - websockets
+- librosa
 - numpy
 
 The script:
@@ -25,12 +26,12 @@ import argparse
 import asyncio
 import json
 
+import librosa
 import numpy as np
 import pybase64 as base64
 import websockets
 
 from vllm.assets.audio import AudioAsset
-from vllm.multimodal.media.audio import load_audio
 
 
 def audio_to_pcm16_base64(audio_path: str) -> str:
@@ -38,7 +39,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
     Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
     """
     # Load audio and resample to 16kHz mono
-    audio, _ = load_audio(audio_path, sr=16000, mono=True)
+    audio, _ = librosa.load(audio_path, sr=16000, mono=True)
     # Convert to PCM16
     pcm16 = (audio * 32767).astype(np.int16)
     # Encode as base64

examples/pooling/token_embed/jina_embeddings_v4_offline.py

@@ -4,74 +4,68 @@
 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,
+)
 
-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,
-    )
+# Create text prompts
+text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+text1_prompt = TextPrompt(prompt=f"Query: {text1}")
 
-    # Create text prompts
-    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+text2 = "浜辺に沈む美しい夕日"
+text2_prompt = TextPrompt(prompt=f"Query: {text2}")
 
-    text2 = "浜辺に沈む美しい夕日"
-    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+# 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},
+)
 
-    # 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)
+# Encode all prompts
+prompts = [text1_prompt, text2_prompt, image_prompt]
+outputs = model.encode(prompts, pooling_task="token_embed")
 
 
-if __name__ == "__main__":
-    main()
+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)

examples/pooling/token_embed/multi_vector_retrieval_offline.py

@@ -4,7 +4,6 @@
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
-from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -14,7 +13,6 @@ 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,
     )
@@ -34,6 +32,15 @@ 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")
 
@@ -43,20 +50,6 @@ 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,11 +7,10 @@ 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 --pooler-config.task token_embed
+vllm serve BAAI/bge-m3
 """
 
 import argparse
-import pprint
 
 import requests
 import torch
@@ -33,8 +32,7 @@ def parse_args():
 
 
 def main(args):
-    pooling_url = f"http://{args.host}:{args.port}/pooling"
-    score_url = f"http://{args.host}:{args.port}/score"
+    api_url = f"http://{args.host}:{args.port}/pooling"
     model_name = args.model
 
     prompts = [
@@ -45,23 +43,11 @@ def main(args):
     ]
     prompt = {"model": model_name, "input": prompts}
 
-    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=api_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()

pyproject.toml

@@ -170,7 +170,6 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
-VALU = "VALU"
 # Walsh-Hadamard Transform
 wht = "wht"
 WHT = "WHT"

requirements/common.txt

@@ -32,7 +32,9 @@ 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.8.post1
-flashinfer-cubin==0.6.8.post1
+flashinfer-python==0.6.7
+flashinfer-cubin==0.6.7
 # 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/rocm.txt

@@ -22,4 +22,4 @@ timm>=1.0.17
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
 # Required for faster safetensors model loading
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2
+fastsafetensors >= 0.2.2

requirements/test/rocm.in

@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
 numba==0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
+fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0

requirements/test/rocm.txt

@@ -76,7 +76,9 @@ attrs==26.1.0
 audioread==3.0.1
     # via librosa
 av==16.1.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/../common.txt
+    #   -r requirements/test/rocm.in
 azure-core==1.39.0
     # via
     #   azure-identity
@@ -275,7 +277,7 @@ fastar==0.10.0
     # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
+fastsafetensors==0.2.2
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
@@ -1331,6 +1333,7 @@ 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

@@ -1085,18 +1085,14 @@ setup(
     install_requires=get_requirements(),
     extras_require={
         # AMD Zen CPU optimizations via zentorch
-        "zen": [
-            "zentorch-weekly==5.2.1.dev20260408"
-        ],  # Zentorch has weekly releases. This pulls the known-good version.
+        "zen": ["zentorch"],
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.2.2"],
         "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

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

@@ -1,98 +0,0 @@
-# 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,10 +9,7 @@ import torch
 import torch.distributed
 
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed.eplb.eplb_communicator import (
-    create_eplb_communicator,
-    has_nixl,
-)
+from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
 from vllm.distributed.eplb.rebalance_execute import (
     move_from_buffer,
     rearrange_expert_weights_inplace,
@@ -361,7 +358,7 @@ def _test_async_transfer_layer_without_mtp_worker(
         communicator.set_stream(cuda_stream)
 
         for layer_idx in range(num_layers):
-            transfer_metadata = asyncio.run(
+            is_unchanged, is_received_locally, recv_metadata = asyncio.run(
                 transfer_layer(
                     old_layer_indices=old_indices_cpu[layer_idx],
                     new_layer_indices=new_indices_cpu[layer_idx],
@@ -376,7 +373,9 @@ def _test_async_transfer_layer_without_mtp_worker(
             move_from_buffer(
                 expert_weights=expert_weights[layer_idx],
                 expert_weights_buffers=expert_buffer,
-                transfer_metadata=transfer_metadata,
+                is_unchanged=is_unchanged,
+                is_received_locally=is_received_locally,
+                recv_metadata=recv_metadata,
                 new_indices=new_indices_cpu[layer_idx].numpy(),
                 ep_rank=ep_rank,
             )
@@ -528,9 +527,7 @@ def _test_rearrange_expert_weights_with_redundancy(
         (4, 8, 8, 16),
     ],
 )
-@pytest.mark.parametrize(
-    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
-)
+@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
 def test_rearrange_expert_weights_with_redundancy(
     world_size,
     num_layers,
@@ -540,8 +537,6 @@ 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(
@@ -638,9 +633,7 @@ 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", "nixl"]
-)
+@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
 def test_async_transfer_layer_without_mtp(
     world_size: int,
     num_layers: int,
@@ -650,8 +643,6 @@ 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], layer)
+    _commit_eplb_maps_for_layer(model_state, new_phy2log, 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[0], layer=0)
+    _commit_eplb_maps_for_layer(model_state, new_phy2log, 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,8 +38,6 @@ 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

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

@@ -1,241 +0,0 @@
-# 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/openai/correctness/test_transcription_api_correctness.py

@@ -13,6 +13,7 @@ import io
 import time
 from statistics import mean, median
 
+import librosa
 import pytest
 import soundfile
 import torch
@@ -20,7 +21,6 @@ from datasets import load_dataset
 from evaluate import load
 from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
 
-from vllm.multimodal.audio import get_audio_duration
 from vllm.tokenizers import get_tokenizer
 
 from ....models.registry import HF_EXAMPLE_MODELS
@@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request):
         trust_remote_code=model_info.trust_remote_code,
     )
 
-    # Warmup call as the first `load_audio` server-side is quite slow.
+    # Warmup call as the first `librosa.load` server-side is quite slow.
     audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
     _ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")
 
@@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time):
 
 def add_duration(sample):
     y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
-    sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
+    sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
     return sample
 
 

tests/entrypoints/openai/realtime/test_realtime_validation.py

@@ -5,6 +5,7 @@ import asyncio
 import json
 import warnings
 
+import librosa
 import numpy as np
 import pybase64 as base64
 import pytest
@@ -13,7 +14,6 @@ import websockets
 from tests.entrypoints.openai.conftest import add_attention_backend
 from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from vllm.assets.audio import AudioAsset
-from vllm.multimodal.media.audio import load_audio
 
 # Increase engine iteration timeout for ROCm where first-use JIT compilation
 # can exceed the default 60s, causing a silent deadlock in feed_tokens.
@@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None:
 def mary_had_lamb_audio_chunks() -> list[str]:
     """Audio split into ~1 second chunks for streaming."""
     path = AudioAsset("mary_had_lamb").get_local_path()
-    audio, _ = load_audio(str(path), sr=16000, mono=True)
+    audio, _ = librosa.load(str(path), sr=16000, mono=True)
 
     # Split into ~0.1 second chunks (1600 samples at 16kHz)
     chunk_size = 1600

tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py

@@ -6,6 +6,7 @@ import asyncio
 import io
 import json
 
+import librosa
 import numpy as np
 import openai
 import pytest
@@ -13,7 +14,6 @@ import pytest_asyncio
 import soundfile as sf
 
 from tests.utils import RemoteOpenAIServer
-from vllm.multimodal.media.audio import load_audio
 from vllm.platforms import current_platform
 
 MODEL_NAME = "openai/whisper-large-v3-turbo"
@@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client):
 @pytest.mark.asyncio
 async def test_long_audio_request(mary_had_lamb, whisper_client):
     mary_had_lamb.seek(0)
-    audio, sr = load_audio(mary_had_lamb)
+    audio, sr = librosa.load(mary_had_lamb)
     # Add small silence after each audio for repeatability in the split process
     audio = np.pad(audio, (0, 1600))
     repeated_audio = np.tile(audio, 10)

tests/entrypoints/openai/speech_to_text/test_translation_validation.py

@@ -7,6 +7,7 @@ import io
 import json
 
 import httpx
+import librosa
 import numpy as np
 import openai
 import pytest
@@ -16,7 +17,6 @@ import soundfile as sf
 from tests.entrypoints.openai.conftest import add_attention_backend
 from tests.utils import RemoteOpenAIServer
 from vllm.logger import init_logger
-from vllm.multimodal.media.audio import load_audio
 
 logger = init_logger(__name__)
 
@@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model):
     if model_name == "google/gemma-3n-E2B-it":
         pytest.skip("Gemma3n does not support long audio requests")
     foscolo.seek(0)
-    audio, sr = load_audio(foscolo)
+    audio, sr = librosa.load(foscolo)
     repeated_audio = np.tile(audio, 2)
     # Repeated audio to buffer
     buffer = io.BytesIO()

tests/entrypoints/openai/test_mm_serde.py

@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
-
-import torch
-
-from vllm.entrypoints.serve.disagg.mm_serde import (
-    decode_mm_kwargs_item,
-    encode_mm_kwargs_item,
-)
-from vllm.entrypoints.serve.disagg.protocol import (
-    MultiModalFeatures,
-    PlaceholderRangeInfo,
-)
-from vllm.multimodal.inputs import (
-    MultiModalBatchedField,
-    MultiModalFieldElem,
-    MultiModalFlatField,
-    MultiModalKwargsItem,
-    MultiModalSharedField,
-)
-
-
-def test_mm_kwargs_item_roundtrip():
-    """Full roundtrip test with all three field types and multiple dtypes."""
-    e1 = MultiModalFieldElem(
-        data=torch.zeros(1000, dtype=torch.bfloat16),
-        field=MultiModalBatchedField(),
-    )
-    e2 = MultiModalFieldElem(
-        data=torch.ones(100, dtype=torch.int32),
-        field=MultiModalSharedField(batch_size=4),
-    )
-    e3 = MultiModalFieldElem(
-        data=torch.randn(20, dtype=torch.float32),
-        field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
-    )
-
-    item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
-    encoded = encode_mm_kwargs_item(item)
-
-    # Encoded result is a base64 string
-    assert isinstance(encoded, str)
-
-    decoded = decode_mm_kwargs_item(encoded)
-
-    assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
-    assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
-    assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
-    assert torch.equal(item["embeds"].data, decoded["embeds"].data)
-    assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
-    assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
-    assert isinstance(decoded["embeds"].field, MultiModalFlatField)
-
-
-def test_mm_kwargs_item_none_data():
-    """Roundtrip with None data field."""
-    elem = MultiModalFieldElem(
-        data=None,
-        field=MultiModalSharedField(batch_size=2),
-    )
-    item = MultiModalKwargsItem({"empty": elem})
-    encoded = encode_mm_kwargs_item(item)
-    decoded = decode_mm_kwargs_item(encoded)
-
-    assert decoded["empty"].data is None
-    assert isinstance(decoded["empty"].field, MultiModalSharedField)
-
-
-def test_mm_kwargs_item_nested_tensors():
-    """Roundtrip with nested tensor data."""
-    nested = [torch.randn(3, 4), torch.randn(5, 4)]
-    elem = MultiModalFieldElem(
-        data=nested,
-        field=MultiModalBatchedField(),
-    )
-    item = MultiModalKwargsItem({"nested": elem})
-    encoded = encode_mm_kwargs_item(item)
-    decoded = decode_mm_kwargs_item(encoded)
-
-    decoded_data = decoded["nested"].data
-    assert len(decoded_data) == 2
-    assert torch.equal(nested[0], decoded_data[0])
-    assert torch.equal(nested[1], decoded_data[1])
-
-
-def test_mm_features_with_kwargs_data():
-    """Test that MultiModalFeatures can carry serialized tensor data."""
-    elem = MultiModalFieldElem(
-        data=torch.randn(5, 3, dtype=torch.float32),
-        field=MultiModalBatchedField(),
-    )
-    item = MultiModalKwargsItem({"pixel_values": elem})
-    encoded = encode_mm_kwargs_item(item)
-
-    features = MultiModalFeatures(
-        mm_hashes={"image": ["abc123"]},
-        mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
-        kwargs_data={"image": [encoded]},
-    )
-
-    # JSON roundtrip
-    json_str = features.model_dump_json()
-    features2 = MultiModalFeatures.model_validate_json(json_str)
-
-    assert features2.mm_hashes == {"image": ["abc123"]}
-    assert features2.kwargs_data is not None
-    assert len(features2.kwargs_data["image"]) == 1
-
-    decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
-    assert torch.equal(elem.data, decoded["pixel_values"].data)

tests/entrypoints/pooling/classify/test_offline.py

@@ -1,12 +1,13 @@
 # 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
+from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask
 
@@ -65,6 +66,18 @@ 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):
@@ -97,12 +110,10 @@ def test_score_api(llm: LLM):
         llm.score("ping", "pong", use_tqdm=False)
 
 
-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "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,7 +436,26 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+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"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -450,11 +469,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_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,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref
 
 import pytest
@@ -37,11 +38,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()
 
 
@@ -73,6 +74,16 @@ 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):
@@ -96,14 +107,10 @@ def test_pooling_params(llm: LLM):
     )
 
 
-@pytest.mark.parametrize(
-    "task", ["token_classify", "classify", "token_embed", "plugin"]
-)
+@pytest.mark.parametrize("task", ["token_classify", "classify", "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,9 +732,28 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize(
-    "task", ["classify", "token_classify", "token_embed", "plugin"]
-)
+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"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -750,8 +769,6 @@ 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,6 +452,25 @@ 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,
@@ -525,7 +544,7 @@ async def test_rerank_max_tokens_per_doc_validation(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     response = requests.post(
         server.url_for("pooling"),
@@ -539,8 +558,6 @@ 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,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref
 
 import pytest
@@ -59,19 +60,22 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "This model does not support the Scoring API."
+    err_msg = "Scoring API is only enabled for num_labels == 1."
     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 == "plugin":
-        err_msg = "No IOProcessor plugin installed."
-    elif task == "classify":
-        err_msg = "Try switching the model's pooling_task via.+"
+    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
     else:
-        err_msg = "Embedding API is not supported by this model.+"
+        if task == "plugin":
+            err_msg = "No IOProcessor plugin installed."
+        else:
+            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", ["classify", "embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,12 +63,9 @@ 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,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref
 
 import pytest
@@ -63,12 +64,15 @@ 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 == "plugin":
-        err_msg = "No IOProcessor plugin installed."
-    elif task == "embed":
-        err_msg = "Try switching the model's pooling_task via.+"
+    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
     else:
-        err_msg = "Classification API is not supported by this model.+"
+        if task == "plugin":
+            err_msg = "No IOProcessor plugin installed."
+        else:
+            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", ["embed", "classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,12 +86,9 @@ 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/entrypoints/serve/disagg/test_serving_multimodal_tokens.py

@@ -1,158 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for multimodal features through the /inference/v1/generate endpoint.
-
-Mirrors test_serving_tokens.py but exercises the multimodal piping
-using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
-/generate -> /detokenize path. Intentionally avoids running the HF
-processor in the pytest parent process to keep os.fork() in sibling
-tests (e.g. test_weight_transfer_llm.py) deadlock-free.
-"""
-
-import os
-
-import httpx
-import pytest
-import pytest_asyncio
-from PIL import Image
-
-from tests.utils import RemoteOpenAIServer
-from vllm.multimodal.utils import encode_image_url
-
-MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
-GEN_ENDPOINT = "/inference/v1/generate"
-RENDER_ENDPOINT = "/v1/chat/completions/render"
-DETOKENIZE_ENDPOINT = "/detokenize"
-
-
-@pytest.fixture(scope="module")
-def test_image():
-    return Image.new("RGB", (224, 224), color=(255, 0, 0))
-
-
-@pytest.fixture(scope="module")
-def server():
-    args = [
-        "--dtype",
-        "bfloat16",
-        "--max-model-len",
-        "4096",
-        "--enforce-eager",
-        "--no-enable-prefix-caching",
-    ]
-
-    envs = os.environ.copy()
-    envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
-
-    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def client(server: RemoteOpenAIServer):
-    transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
-    headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
-    async with httpx.AsyncClient(
-        transport=transport,
-        base_url=server.url_root,
-        timeout=600,
-        headers=headers,
-    ) as c:
-        yield c
-
-
-@pytest.mark.asyncio
-async def test_render_to_generate_roundtrip(client, test_image):
-    """End-to-end: render a multimodal chat -> feed into generate -> decode.
-
-    All preprocessing and detokenization happens in the server subprocess;
-    the pytest parent never imports transformers or touches torch tensors.
-    """
-    data_url = encode_image_url(test_image, format="PNG")
-
-    render_payload = {
-        "model": MODEL_NAME,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": data_url}},
-                    {
-                        "type": "text",
-                        "text": "What color is this image? Answer in one word.",
-                    },
-                ],
-            }
-        ],
-    }
-
-    render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
-    render_resp.raise_for_status()
-    render_data = render_resp.json()
-
-    # Validate render output structure: keys exist and values are non-empty
-    # and well-typed.
-    assert "token_ids" in render_data
-    assert isinstance(render_data["token_ids"], list)
-    assert len(render_data["token_ids"]) > 0
-    assert all(isinstance(t, int) for t in render_data["token_ids"])
-
-    assert "features" in render_data
-    features = render_data["features"]
-    assert features is not None
-    assert isinstance(features, dict)
-
-    assert "mm_hashes" in features
-    assert "image" in features["mm_hashes"]
-    image_hashes = features["mm_hashes"]["image"]
-    assert isinstance(image_hashes, list)
-    assert len(image_hashes) > 0
-    assert all(isinstance(h, str) and h for h in image_hashes)
-
-    assert "mm_placeholders" in features
-    assert "image" in features["mm_placeholders"]
-    image_placeholders = features["mm_placeholders"]["image"]
-    assert isinstance(image_placeholders, list)
-    assert len(image_placeholders) > 0
-    for p in image_placeholders:
-        assert isinstance(p.get("offset"), int)
-        assert isinstance(p.get("length"), int)
-        assert p["length"] > 0
-
-    assert "kwargs_data" in features
-    assert "image" in features["kwargs_data"]
-    assert len(features["kwargs_data"]["image"]) > 0
-
-    # Build generate request from render output
-    generate_payload = render_data
-    generate_payload["sampling_params"] = {
-        "max_tokens": 10,
-        "temperature": 0.0,
-    }
-
-    gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
-    gen_resp.raise_for_status()
-    gen_data = gen_resp.json()
-
-    assert "choices" in gen_data
-    assert isinstance(gen_data["choices"], list)
-    assert len(gen_data["choices"]) >= 1
-    choice = gen_data["choices"][0]
-    assert "token_ids" in choice
-    assert isinstance(choice["token_ids"], list)
-    assert len(choice["token_ids"]) > 0
-    assert all(isinstance(t, int) for t in choice["token_ids"])
-
-    detok_resp = await client.post(
-        DETOKENIZE_ENDPOINT,
-        json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
-    )
-    detok_resp.raise_for_status()
-    detok_data = detok_resp.json()
-    assert "prompt" in detok_data
-    text = detok_data["prompt"]
-    assert isinstance(text, str)
-    assert len(text) > 0
-    assert "red" in text.lower(), (
-        f"Expected model to identify the red image, got: {text!r}"
-    )

tests/ir/ir_test_utils.py

@@ -1,46 +0,0 @@
-# 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,68 +495,3 @@ 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,12 +4,9 @@
 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])
@@ -28,35 +25,33 @@ 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=DEVICE_TYPE
+        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
     )
     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=DEVICE_TYPE).view(
-        1, 1, -1
-    )
+    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").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=DEVICE_TYPE)
+    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
 
     # 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=DEVICE_TYPE)
-    v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
+    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")
 
     # o will have the same shape as q
-    o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
+    o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
 
-    lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
+    lse = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
 
-    b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
+    b_seq_len = torch.full((B,), seq_len, device="cuda")
 
     attn_logits = torch.empty(
         (B, H_Q, num_kv_splits, D_V + 1),
         dtype=torch.float32,
-        device=DEVICE_TYPE,
+        device="cuda",
     )
 
     # Call the original implementation.
@@ -132,27 +127,25 @@ 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=DEVICE_TYPE
+        0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
     )
     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=DEVICE_TYPE).view(
-        1, 1, -1
-    )
+    req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").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=DEVICE_TYPE)
+    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
 
     # Create BF16 K/V as reference
-    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)
+    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")
 
     # --- BF16 reference ---
-    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)
+    o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
+    lse_ref = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
     attn_logits = torch.empty(
-        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
+        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
     )
 
     if PAGE_SIZE == 1:
@@ -163,7 +156,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=DEVICE_TYPE),
+            b_seq_len=torch.full((B,), seq_len, device="cuda"),
             attn_logits=attn_logits,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -178,7 +171,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=DEVICE_TYPE),
+            b_seq_len=torch.full((B,), seq_len, device="cuda"),
             attn_logits=attn_logits,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -189,10 +182,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=DEVICE_TYPE)
-    lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
+    o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
+    lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
     attn_logits_fp8 = torch.empty(
-        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
+        (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
     )
 
     if PAGE_SIZE == 1:
@@ -203,7 +196,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=DEVICE_TYPE),
+            b_seq_len=torch.full((B,), seq_len, device="cuda"),
             attn_logits=attn_logits_fp8,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,
@@ -220,7 +213,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=DEVICE_TYPE),
+            b_seq_len=torch.full((B,), seq_len, device="cuda"),
             attn_logits=attn_logits_fp8,
             num_kv_splits=num_kv_splits,
             sm_scale=sm_scale,

tests/kernels/attention/test_triton_prefill_attention.py

@@ -5,11 +5,8 @@ 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,
@@ -95,19 +92,17 @@ 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=DEVICE_TYPE
-    )
+    seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
     total_tokens = seq_lens.sum().item()
 
     # Create batch start locations
-    b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
+    b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
     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=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)
+    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")
     o = torch.zeros_like(q)
 
     # Call Triton kernel
@@ -174,19 +169,17 @@ 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=DEVICE_TYPE
-    )
+    seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
     total_tokens = seq_lens.sum().item()
 
     # Create batch start locations
-    b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
+    b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
     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=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)
+    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")
     o = torch.zeros_like(q)
 
     # Call Triton kernel

tests/kernels/attention/test_triton_unified_attention.py

@@ -10,8 +10,6 @@ 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]
@@ -116,7 +114,7 @@ def test_triton_unified_attn(
     q_dtype: torch.dtype | None,
     seq_threshold_3D: int,
 ) -> None:
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
 
     set_random_seed(0)
     num_seqs = len(seq_lens)
@@ -251,7 +249,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(DEVICE_TYPE)
+    torch.set_default_device("cuda")
 
     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_family",
+    "vllm.utils.flashinfer.current_platform.is_device_capability",
     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_family",
+    "vllm.utils.flashinfer.current_platform.is_device_capability",
     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_family",
+    "vllm.utils.flashinfer.current_platform.is_device_capability",
     return_value=True,
 )
 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)

tests/kernels/ir/test_ir_ops.py

@@ -1,20 +0,0 @@
-# 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", NUM_TOKENS)
-@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
+@pytest.mark.parametrize("n_tokens", [1, 8, 17])
+@pytest.mark.parametrize("hidden_size", [16, 4096, 8192])
 @pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
@@ -53,9 +53,7 @@ class TestRMSNorm:
         torch.set_default_device(current_platform.device_type)
 
     def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
-        x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
-            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
-        )
+        x, weight = rms_norm_inputs(4, 8, dtype)
         out = rms_norm_native(x, weight, epsilon=epsilon)
 
         # Check shape, dtype, device
@@ -73,59 +71,59 @@ class TestRMSNorm:
         out4 = rms_norm_native(x, None, epsilon=epsilon)
         torch.testing.assert_close(out3, out4)
 
-    @pytest.mark.parametrize("provider", supported_providers(ir.ops.rms_norm))
+    @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels"])
     def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
         impl = ir.ops.rms_norm.impls[provider]
-        x, weight, eps = ir.ops.rms_norm.generate_inputs(
-            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
+        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
         )
-        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_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)
+            out_impl2 = ir.ops.rms_norm(*args)
+
+        # exact match
+        torch.testing.assert_close(out_impl2, out_impl, rtol=0.0, atol=0.0)
 
         # none of these support variance_size override
-        assert not impl.supports_args(x, weight, eps, 4)
-        assert not impl.supports_args(x, weight, eps, variance_size=4)
+        assert not impl.supports_args(x, weight, epsilon, 4)
+        assert not impl.supports_args(x, weight, epsilon, variance_size=4)
 
         # test weight=None behavior
-        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)
+        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,
+        )
 
     @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")
 
-        args = ir.ops.rms_norm.generate_inputs(
-            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
-        )
+        x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
+        args = (x, weight, epsilon, None)
 
         # 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,70 +135,3 @@ 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_gemma4router.py

@@ -1,57 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import pytest
-import torch
-
-from vllm.model_executor.models.gemma4 import (
-    gemma4_fused_routing_kernel_triton,
-    gemma4_routing_function_torch,
-)
-
-
-def sort_by_id(w, ids):
-    order = ids.argsort(dim=-1)
-    return w.gather(1, order), ids.gather(1, order)
-
-
-# Gemma4 MoE Model has context length of 250K
-# the minus 1 is to ensure that edge cases are tested
-@pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000])
-@pytest.mark.parametrize("num_experts", [128])  # gemma4 moe experts
-@pytest.mark.parametrize("topk", [8])  # gemma4 topk
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
-def test_gemma4_routing_kernel_triton(
-    num_tokens: int,
-    num_experts: int,
-    topk: int,
-    dtype: torch.dtype,
-):
-    torch.manual_seed(0)
-
-    gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda")
-    scales = torch.rand(num_experts, dtype=torch.float32, device="cuda")
-
-    ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales)
-    tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales)
-
-    # Sort by expert id — to remove tie-breaking differences
-    ref_ws, ref_is = sort_by_id(ref_w, ref_ids)
-    tri_ws, tri_is = sort_by_id(tri_w, tri_ids)
-
-    ids_match = (ref_is == tri_is).all().item()
-    weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2)
-    all_match = ids_match and weights_match
-    max_err = (ref_ws - tri_ws).abs().max().item()
-    print(
-        f"T={num_tokens:5d} E={num_experts:4d} K={topk} "
-        f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}"
-    )
-    if not all_match:
-        bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0]
-        if len(bad):
-            r = bad[0].item()
-            print(
-                f"  first bad row {r}: ref_ids={ref_ids[r].tolist()} "
-                f"tri_ids={tri_ids[r].tolist()}"
-            )
-        assert all_match