updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

.buildkite/nightly-benchmarks/scripts/convert-results-json-to-markdown.py

@@ -368,7 +368,7 @@ if __name__ == "__main__":
         # The GPUs sometimes come in format of "GPUTYPE\nGPUTYPE\n...",
         # we want to turn it into "8xGPUTYPE"
         df["GPU"] = df["GPU"].apply(
-            lambda x: f"{len(x.splitlines())}x{x.splitlines()[0]}"
+            lambda x: f"{len(x.split('\n'))}x{x.split('\n')[0]}"
         )
 
     # get markdown tables

.buildkite/pyproject.toml

@@ -0,0 +1,46 @@
+# This local pyproject file is part of the migration from yapf to ruff format.
+# It uses the same core rules as the main pyproject.toml file, but with the
+# following differences:
+# - ruff line length is overridden to 88
+# - deprecated typing ignores (UP006, UP035) have been removed
+
+[tool.ruff]
+line-length = 88
+
+[tool.ruff.lint.per-file-ignores]
+"vllm/third_party/**" = ["ALL"]
+"vllm/version.py" = ["F401"]
+"vllm/_version.py" = ["ALL"]
+
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    "I",
+    # flake8-logging-format
+    "G",
+]
+ignore = [
+    # star imports
+    "F405", "F403",
+    # lambda expression assignment
+    "E731",
+    # Loop control variable not used within loop body
+    "B007",
+    # f-string format
+    "UP032",
+    # Can remove once 3.10+ is the minimum Python version
+    "UP007",
+]
+
+[tool.ruff.format]
+docstring-code-format = true

.buildkite/release-pipeline.yaml

@@ -150,16 +150,11 @@ steps:
       queue: cpu_queue_postmerge
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64"
-      - "docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64"
-      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 vllm/vllm-openai:nightly-x86_64"
-      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 vllm/vllm-openai:nightly-aarch64"
-      - "docker push vllm/vllm-openai:nightly-x86_64"
-      - "docker push vllm/vllm-openai:nightly-aarch64"
-      - "docker manifest create vllm/vllm-openai:nightly vllm/vllm-openai:nightly-x86_64 vllm/vllm-openai:nightly-aarch64 --amend"
-      - "docker manifest create vllm/vllm-openai:nightly-$BUILDKITE_COMMIT vllm/vllm-openai:nightly-x86_64 vllm/vllm-openai:nightly-aarch64 --amend"
-      - "docker manifest push vllm/vllm-openai:nightly"
-      - "docker manifest push vllm/vllm-openai:nightly-$BUILDKITE_COMMIT"
+      - "docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT vllm/vllm-openai:nightly"
+      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT vllm/vllm-openai:nightly-$BUILDKITE_COMMIT"
+      - "docker push vllm/vllm-openai:nightly"
+      - "docker push vllm/vllm-openai:nightly-$BUILDKITE_COMMIT"
       # Clean up old nightly builds (keep only last 14)
       - "bash .buildkite/scripts/cleanup-nightly-builds.sh"
     plugins:
@@ -168,4 +163,3 @@ steps:
           password-env: DOCKERHUB_TOKEN
     env:
       DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"

.buildkite/scripts/cleanup-nightly-builds.sh

@@ -8,41 +8,20 @@ set -ex
 # DockerHub API endpoint for vllm/vllm-openai repository
 REPO_API_URL="https://hub.docker.com/v2/repositories/vllm/vllm-openai/tags"
 
-# Get DockerHub credentials from environment
+# Get DockerHub token from environment
 if [ -z "$DOCKERHUB_TOKEN" ]; then
     echo "Error: DOCKERHUB_TOKEN environment variable is not set"
     exit 1
 fi
 
-if [ -z "$DOCKERHUB_USERNAME" ]; then
-    echo "Error: DOCKERHUB_USERNAME environment variable is not set"
-    exit 1
-fi
-
-# Get DockerHub bearer token
-echo "Getting DockerHub bearer token..."
-set +x
-BEARER_TOKEN=$(curl -s -X POST \
-    -H "Content-Type: application/json" \
-    -d "{\"username\": \"$DOCKERHUB_USERNAME\", \"password\": \"$DOCKERHUB_TOKEN\"}" \
-    "https://hub.docker.com/v2/users/login" | jq -r '.token')
-set -x
-
-if [ -z "$BEARER_TOKEN" ] || [ "$BEARER_TOKEN" = "null" ]; then
-    echo "Error: Failed to get DockerHub bearer token"
-    exit 1
-fi
-
 # Function to get all tags from DockerHub
 get_all_tags() {
     local page=1
     local all_tags=""
     
     while true; do
-        set +x
-        local response=$(curl -s -H "Authorization: Bearer $BEARER_TOKEN" \
+        local response=$(curl -s -H "Authorization: Bearer $DOCKERHUB_TOKEN" \
             "$REPO_API_URL?page=$page&page_size=100")
-        set -x
         
         # Get both last_updated timestamp and tag name, separated by |
         local tags=$(echo "$response" | jq -r '.results[] | select(.name | startswith("nightly-")) | "\(.last_updated)|\(.name)"')
@@ -64,9 +43,7 @@ delete_tag() {
     echo "Deleting tag: $tag_name"
     
     local delete_url="https://hub.docker.com/v2/repositories/vllm/vllm-openai/tags/$tag_name"
-    set +x
-    local response=$(curl -s -X DELETE -H "Authorization: Bearer $BEARER_TOKEN" "$delete_url")
-    set -x
+    local response=$(curl -s -X DELETE -H "Authorization: Bearer $DOCKERHUB_TOKEN" "$delete_url")
     
     if echo "$response" | jq -e '.detail' > /dev/null 2>&1; then
         echo "Warning: Failed to delete tag $tag_name: $(echo "$response" | jq -r '.detail')"

.buildkite/test-pipeline.yaml

@@ -397,7 +397,6 @@ steps:
     - pytest -v -s compile/test_pass_manager.py
     - pytest -v -s compile/test_fusion.py
     - pytest -v -s compile/test_fusion_attn.py
-    - pytest -v -s compile/test_functionalization.py
     - pytest -v -s compile/test_silu_mul_quant_fusion.py
     - pytest -v -s compile/test_sequence_parallelism.py
     - pytest -v -s compile/test_async_tp.py
@@ -477,7 +476,6 @@ steps:
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
-  - vllm/model_executor/layers/mamba/ops
   commands:
     - pytest -v -s kernels/mamba
 
@@ -835,11 +833,11 @@ steps:
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
 
-- label: Blackwell GPT-OSS Eval
+- label: GPT-OSS Eval (Blackwell)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"
   gpu: b200
-  optional: true # run on nightlies
+  optional: true # disable while debugging
   source_file_dependencies:
   - tests/evals/gpt_oss
   - vllm/model_executor/models/gpt_oss.py
@@ -866,16 +864,6 @@ steps:
   commands:
     - pytest -s -v tests/quantization/test_blackwell_moe.py
 
-- label: Blackwell LM Eval Small Models
-  timeout_in_minutes: 75
-  gpu: b200
-  optional: true # run on nightlies
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt --tp-size=1
-
 #####  1 GPU test  #####
 #####  multi gpus test  #####
 

.github/CODEOWNERS

@@ -23,7 +23,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
 /vllm/config @simon-mo @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
-/vllm/config/cache.py @simon-mo @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg @heheda12345
 
 # vLLM V1
 /vllm/v1 @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat

.github/workflows/stale.yml

@@ -13,7 +13,7 @@ jobs:
       actions: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/stale@5f858e3efba33a5ca4407a664cc011ad407f2008 # v10.1.0
+      - uses: actions/stale@3a9db7e6a41a89f618792c92c0e97cc736e1b13f # v10.0.0
         with:
           # Increasing this value ensures that changes to this workflow
           # propagate to all issues and PRs in days rather than months

.pre-commit-config.yaml

@@ -6,16 +6,28 @@ default_stages:
   - manual # Run in CI
 exclude: 'vllm/third_party/.*'
 repos:
-- repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.13.3
+- repo: https://github.com/google/yapf
+  rev: v0.43.0
   hooks:
-  - id: ruff-check
+  - id: yapf
+    args: [--in-place, --verbose]
+    # Keep the same list from yapfignore here to avoid yapf failing without any inputs
+    exclude: '(.buildkite|benchmarks|build|examples)/.*'
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.11.7
+  hooks:
+  - id: ruff
     args: [--output-format, github, --fix]
   - id: ruff-format
+    files: ^(.buildkite|benchmarks|examples)/.*
 - repo: https://github.com/crate-ci/typos
   rev: v1.35.5
   hooks:
   - id: typos
+- repo: https://github.com/PyCQA/isort
+  rev: 6.0.1
+  hooks:
+  - id: isort
 - repo: https://github.com/pre-commit/mirrors-clang-format
   rev: v20.1.3
   hooks:

benchmarks/benchmark_block_pool.py

@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
 
-from benchmark_utils import TimeCollector
 from tabulate import tabulate
 
+from benchmark_utils import TimeCollector
 from vllm.utils import FlexibleArgumentParser
 from vllm.v1.core.block_pool import BlockPool
 

benchmarks/benchmark_ngram_proposer.py

@@ -5,9 +5,9 @@ import time
 from unittest import mock
 
 import numpy as np
-from benchmark_utils import TimeCollector
 from tabulate import tabulate
 
+from benchmark_utils import TimeCollector
 from vllm.config import (
     CacheConfig,
     DeviceConfig,
@@ -164,7 +164,7 @@ def invoke_main() -> None:
     )
     parser.add_argument(
         "--batched", action="store_true", help="consider time to prepare batch"
-    )
+    )  # noqa: E501
     parser.add_argument(
         "--num-iteration",
         type=int,

benchmarks/benchmark_serving_structured_output.py

@@ -37,13 +37,14 @@ from typing import Optional
 import datasets
 import numpy as np
 import pandas as pd
+from tqdm.asyncio import tqdm
+from transformers import PreTrainedTokenizerBase
+
 from backend_request_func import (
     ASYNC_REQUEST_FUNCS,
     RequestFuncInput,
     RequestFuncOutput,
 )
-from tqdm.asyncio import tqdm
-from transformers import PreTrainedTokenizerBase
 
 try:
     from vllm.transformers_utils.tokenizer import get_tokenizer
@@ -909,13 +910,13 @@ def create_argument_parser():
     parser.add_argument(
         "--tokenizer",
         type=str,
-        help="Name or path of the tokenizer, if not using the default tokenizer.",
+        help="Name or path of the tokenizer, if not using the default tokenizer.",  # noqa: E501
     )
     parser.add_argument(
         "--tokenizer-mode",
         type=str,
         default="auto",
-        help="Name or path of the tokenizer, if not using the default tokenizer.",
+        help="Name or path of the tokenizer, if not using the default tokenizer.",  # noqa: E501
     )
     parser.add_argument(
         "--num-prompts",

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# fmt: off
 # ruff: noqa: E501
 import time
 
@@ -19,21 +20,19 @@ from vllm.utils.deep_gemm import (
 )
 
 
-def benchmark_shape(
-    m: int,
-    n: int,
-    k: int,
-    warmup: int = 100,
-    repeat: int = 10000,
-    verbose: bool = False,
-) -> dict:
+def benchmark_shape(m: int,
+                    n: int,
+                    k: int,
+                    warmup: int = 100,
+                    repeat: int = 10000,
+                    verbose: bool = False) -> dict:
     """Benchmark all implementations for a specific (m, n, k) shape."""
     if verbose:
         print(f"\n=== Benchmarking shape: m={m}, n={n}, k={k} ===")
 
     # Create test tensors
-    A = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
-    B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
+    A = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
+    B = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
 
     # Reference result in BF16
     torch.cuda.synchronize()
@@ -50,39 +49,34 @@ def benchmark_shape(
     # Pre-quantize A for all implementations
     A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(A, block_size[1])
     A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
-    C_deepgemm = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+    C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
     A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
     A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
-        A, block_size[1], column_major_scales=True
-    )
+        A, block_size[1], column_major_scales=True)
 
     # === DeepGEMM Implementation ===
     def deepgemm_gemm():
-        fp8_gemm_nt(
-            (A_deepgemm, A_scale_deepgemm), (B_deepgemm, B_scale_deepgemm), C_deepgemm
-        )
+        fp8_gemm_nt((A_deepgemm, A_scale_deepgemm),
+                                       (B_deepgemm, B_scale_deepgemm),
+                                       C_deepgemm)
         return C_deepgemm
 
     # === vLLM Triton Implementation ===
     def vllm_triton_gemm():
-        return w8a8_triton_block_scaled_mm(
-            A_vllm,
-            B_vllm,
-            A_scale_vllm,
-            B_scale_vllm,
-            block_size,
-            output_dtype=torch.bfloat16,
-        )
+        return w8a8_triton_block_scaled_mm(A_vllm,
+                                     B_vllm,
+                                     A_scale_vllm,
+                                     B_scale_vllm,
+                                     block_size,
+                                     output_dtype=torch.bfloat16)
 
     # === vLLM CUTLASS Implementation ===
     def vllm_cutlass_gemm():
-        return ops.cutlass_scaled_mm(
-            A_vllm_cutlass,
-            B_vllm.T,
-            scale_a=A_scale_vllm_cutlass,
-            scale_b=B_scale_vllm.T,
-            out_dtype=torch.bfloat16,
-        )
+        return ops.cutlass_scaled_mm(A_vllm_cutlass,
+                                     B_vllm.T,
+                                     scale_a=A_scale_vllm_cutlass,
+                                     scale_b=B_scale_vllm.T,
+                                     out_dtype=torch.bfloat16)
 
     # Run correctness check first
     if verbose:
@@ -99,23 +93,26 @@ def benchmark_shape(
         print(f"DeepGEMM vs Reference difference: {deepgemm_diff:.6f}")
         print(f"vLLM Triton vs Reference difference: {vllm_triton_diff:.6f}")
         print(f"vLLM CUTLASS vs Reference difference: {vllm_cutlass_diff:.6f}")
-        print(
-            "vLLM Triton vs DeepGEMM difference: "
-            f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}"
-        )
-        print(
-            "vLLM CUTLASS vs DeepGEMM difference: "
-            f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}"
-        )
+        print("vLLM Triton vs DeepGEMM difference: "
+              f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}")
+        print("vLLM CUTLASS vs DeepGEMM difference: "
+              f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}")
 
     # Benchmark implementations
     implementations = {
         "DeepGEMM": deepgemm_gemm,
         "vLLM Triton": vllm_triton_gemm,
-        "vLLM CUTLASS": vllm_cutlass_gemm,
+        "vLLM CUTLASS": vllm_cutlass_gemm
     }
 
-    benchmark_results = {"shape": {"m": m, "n": n, "k": k}, "implementations": {}}
+    benchmark_results = {
+        "shape": {
+            "m": m,
+            "n": n,
+            "k": k
+        },
+        "implementations": {}
+    }
 
     for name, func in implementations.items():
         # Warmup
@@ -143,36 +140,38 @@ def benchmark_shape(
             "tflops": tflops,
             "gb_s": gb_s,
             "diff": {
-                "DeepGEMM": 0.0
-                if name == "DeepGEMM"
-                else calc_diff(func(), C_deepgemm),
-                "Reference": deepgemm_diff
-                if name == "DeepGEMM"
-                else (vllm_triton_diff if name == "vLLM Triton" else vllm_cutlass_diff),
-            },
+                "DeepGEMM":
+                0.0 if name == "DeepGEMM" else calc_diff(func(), C_deepgemm),
+                "Reference":
+                deepgemm_diff if name == "DeepGEMM" else
+                (vllm_triton_diff
+                 if name == "vLLM Triton" else vllm_cutlass_diff)
+            }
         }
 
         if verbose:
-            print(f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s")
+            print(
+                f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s"
+            )
 
     # Calculate speedups
     baseline = benchmark_results["implementations"]["DeepGEMM"]["time_ms"]
     for name, data in benchmark_results["implementations"].items():
         if name != "DeepGEMM":
             speedup = baseline / data["time_ms"]
-            benchmark_results["implementations"][name]["speedup_vs_deepgemm"] = speedup
+            benchmark_results["implementations"][name][
+                "speedup_vs_deepgemm"] = speedup
             if verbose:
-                print(
-                    f"DeepGEMM is {1 / speedup:.2f}x "
-                    f"{'faster' if 1 / speedup > 1 else 'slower'} than {name}"
-                )
+                print(f"DeepGEMM is {1/speedup:.2f}x "
+                      f"{'faster' if 1/speedup > 1 else 'slower'} than {name}")
 
-    vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"]["time_ms"]
-    vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"]["time_ms"]
+    vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"][
+        "time_ms"]
+    vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"][
+        "time_ms"]
     cutlass_vs_triton = vllm_triton_time / vllm_cutlass_time
-    benchmark_results["implementations"]["vLLM CUTLASS"]["speedup_vs_triton"] = (
-        cutlass_vs_triton
-    )
+    benchmark_results["implementations"]["vLLM CUTLASS"][
+        "speedup_vs_triton"] = cutlass_vs_triton
     if verbose:
         print(
             f"vLLM CUTLASS is {cutlass_vs_triton:.2f}x "
@@ -184,7 +183,8 @@ def benchmark_shape(
 
 def format_table_row(values, widths):
     """Format a row with specified column widths."""
-    return "| " + " | ".join(f"{val:{w}}" for val, w in zip(values, widths)) + " |"
+    return "| " + " | ".join(f"{val:{w}}"
+                             for val, w in zip(values, widths)) + " |"
 
 
 def print_table(headers, rows, title=None):
@@ -292,50 +292,38 @@ def run_benchmarks(verbose: bool = False):
     for result in all_results:
         shape = result["shape"]
         impl_data = result["implementations"]["DeepGEMM"]
-        deepgemm_rows.append(
-            [
-                shape["m"],
-                shape["n"],
-                shape["k"],
-                f"{impl_data['time_us']:.1f}",
-                f"{impl_data['tflops']:.1f}",
-                f"{impl_data['gb_s']:.1f}",
-            ]
-        )
+        deepgemm_rows.append([
+            shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
+            f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}"
+        ])
 
-    print_table(deepgemm_headers, deepgemm_rows, title="DeepGEMM Implementation:")
+    print_table(deepgemm_headers,
+                deepgemm_rows,
+                title="DeepGEMM Implementation:")
 
     # Print vLLM Triton table
-    triton_headers = ["m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"]
+    triton_headers = [
+        "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"
+    ]
     triton_rows = []
     for result in all_results:
         shape = result["shape"]
         impl_data = result["implementations"]["vLLM Triton"]
         speedup = impl_data.get("speedup_vs_deepgemm", 1.0)
-        triton_rows.append(
-            [
-                shape["m"],
-                shape["n"],
-                shape["k"],
-                f"{impl_data['time_us']:.1f}",
-                f"{impl_data['tflops']:.1f}",
-                f"{impl_data['gb_s']:.1f}",
-                format_speedup(speedup),
-            ]
-        )
+        triton_rows.append([
+            shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
+            f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
+            format_speedup(speedup)
+        ])
 
-    print_table(triton_headers, triton_rows, title="vLLM Triton Implementation:")
+    print_table(triton_headers,
+                triton_rows,
+                title="vLLM Triton Implementation:")
 
     # Print vLLM CUTLASS table
     cutlass_headers = [
-        "m",
-        "n",
-        "k",
-        "Time (μs)",
-        "TFLOPS",
-        "GB/s",
-        "vs DeepGEMM",
-        "vs Triton",
+        "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM",
+        "vs Triton"
     ]
     cutlass_rows = []
     for result in all_results:
@@ -343,27 +331,28 @@ def run_benchmarks(verbose: bool = False):
         impl_data = result["implementations"]["vLLM CUTLASS"]
         vs_deepgemm = impl_data.get("speedup_vs_deepgemm", 1.0)
         vs_triton = impl_data.get("speedup_vs_triton", 1.0)
-        cutlass_rows.append(
-            [
-                shape["m"],
-                shape["n"],
-                shape["k"],
-                f"{impl_data['time_us']:.1f}",
-                f"{impl_data['tflops']:.1f}",
-                f"{impl_data['gb_s']:.1f}",
-                format_speedup(vs_deepgemm),
-                format_speedup(vs_triton),
-            ]
-        )
+        cutlass_rows.append([
+            shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
+            f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
+            format_speedup(vs_deepgemm),
+            format_speedup(vs_triton)
+        ])
 
-    print_table(cutlass_headers, cutlass_rows, title="vLLM CUTLASS Implementation:")
+    print_table(cutlass_headers,
+                cutlass_rows,
+                title="vLLM CUTLASS Implementation:")
 
     # Calculate and print averages
     print("\n===== AVERAGE PERFORMANCE =====")
 
     implementations = ["DeepGEMM", "vLLM Triton", "vLLM CUTLASS"]
     avg_metrics = {
-        impl: {"tflops": 0, "gb_s": 0, "time_ms": 0} for impl in implementations
+        impl: {
+            "tflops": 0,
+            "gb_s": 0,
+            "time_ms": 0
+        }
+        for impl in implementations
     }
 
     for result in all_results:
@@ -381,9 +370,9 @@ def run_benchmarks(verbose: bool = False):
         avg_tflops = avg_metrics[impl]["tflops"] / num_shapes
         avg_mem_bw = avg_metrics[impl]["gb_s"] / num_shapes
         avg_time = avg_metrics[impl]["time_ms"] / num_shapes
-        avg_rows.append(
-            [impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"]
-        )
+        avg_rows.append([
+            impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"
+        ])
 
     print_table(avg_headers, avg_rows)
 
@@ -391,19 +380,21 @@ def run_benchmarks(verbose: bool = False):
     avg_speedups = {
         "DeepGEMM vs vLLM Triton": 0,
         "DeepGEMM vs vLLM CUTLASS": 0,
-        "vLLM CUTLASS vs vLLM Triton": 0,
+        "vLLM CUTLASS vs vLLM Triton": 0
     }
 
     for result in all_results:
         deepgemm_time = result["implementations"]["DeepGEMM"]["time_ms"]
         vllm_triton_time = result["implementations"]["vLLM Triton"]["time_ms"]
-        vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"]["time_ms"]
+        vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"][
+            "time_ms"]
 
-        avg_speedups["DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
-        avg_speedups["DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
-        avg_speedups["vLLM CUTLASS vs vLLM Triton"] += (
-            vllm_triton_time / vllm_cutlass_time
-        )
+        avg_speedups[
+            "DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
+        avg_speedups[
+            "DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
+        avg_speedups[
+            "vLLM CUTLASS vs vLLM Triton"] += vllm_triton_time / vllm_cutlass_time
 
     print("\n===== AVERAGE SPEEDUPS =====")
     speedup_headers = ["Comparison", "Speedup"]
@@ -421,7 +412,8 @@ def run_benchmarks(verbose: bool = False):
 
     for result in all_results:
         for impl in implementations:
-            avg_diff[impl] += result["implementations"][impl]["diff"]["Reference"]
+            avg_diff[impl] += result["implementations"][impl]["diff"][
+                "Reference"]
 
     diff_headers = ["Implementation", "Avg Diff vs Reference"]
     diff_rows = []

benchmarks/pyproject.toml

@@ -0,0 +1,49 @@
+# This local pyproject file is part of the migration from yapf to ruff format.
+# It uses the same core rules as the main pyproject.toml file, but with the
+# following differences:
+# - ruff line length is overridden to 88
+# - deprecated typing ignores (UP006, UP035) have been removed
+
+[tool.ruff]
+line-length = 88
+
+[tool.ruff.lint.per-file-ignores]
+"vllm/third_party/**" = ["ALL"]
+"vllm/version.py" = ["F401"]
+"vllm/_version.py" = ["ALL"]
+
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    "I",
+    # flake8-logging-format
+    "G",
+]
+ignore = [
+    # star imports
+    "F405", "F403",
+    # lambda expression assignment
+    "E731",
+    # Loop control variable not used within loop body
+    "B007",
+    # f-string format
+    "UP032",
+    # Can remove once 3.10+ is the minimum Python version
+    "UP007",
+]
+
+[tool.ruff.lint.isort]
+known-first-party = ["vllm"]
+
+[tool.ruff.format]
+docstring-code-format = true

cmake/cpu_extension.cmake

@@ -213,7 +213,6 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
         endif()
         set(ONEDNN_AARCH64_USE_ACL "ON")
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
-        add_compile_definitions(VLLM_USE_ACL)
     endif()
 
     set(ONEDNN_LIBRARY_TYPE "STATIC")
@@ -227,7 +226,7 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
     set(ONEDNN_ENABLE_ITT_TASKS "OFF")
     set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
     set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
-    set(ONEDNN_VERBOSE "ON")
+    set(ONEDNN_VERBOSE "OFF")
     set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
 
     FetchContent_MakeAvailable(oneDNN)

cmake/hipify.py

@@ -16,7 +16,7 @@ import shutil
 
 from torch.utils.hipify.hipify_python import hipify
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     parser = argparse.ArgumentParser()
 
     # Project directory where all the source + include files live.
@@ -34,14 +34,15 @@ if __name__ == "__main__":
     )
 
     # Source files to convert.
-    parser.add_argument(
-        "sources", help="Source files to hipify.", nargs="*", default=[]
-    )
+    parser.add_argument("sources",
+                        help="Source files to hipify.",
+                        nargs="*",
+                        default=[])
 
     args = parser.parse_args()
 
     # Limit include scope to project_dir only
-    includes = [os.path.join(args.project_dir, "*")]
+    includes = [os.path.join(args.project_dir, '*')]
 
     # Get absolute path for all source files.
     extra_files = [os.path.abspath(s) for s in args.sources]
@@ -50,31 +51,25 @@ if __name__ == "__main__":
     # The directory might already exist to hold object files so we ignore that.
     shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True)
 
-    hipify_result = hipify(
-        project_directory=args.project_dir,
-        output_directory=args.output_dir,
-        header_include_dirs=[],
-        includes=includes,
-        extra_files=extra_files,
-        show_detailed=True,
-        is_pytorch_extension=True,
-        hipify_extra_files_only=True,
-    )
+    hipify_result = hipify(project_directory=args.project_dir,
+                           output_directory=args.output_dir,
+                           header_include_dirs=[],
+                           includes=includes,
+                           extra_files=extra_files,
+                           show_detailed=True,
+                           is_pytorch_extension=True,
+                           hipify_extra_files_only=True)
 
     hipified_sources = []
     for source in args.sources:
         s_abs = os.path.abspath(source)
-        hipified_s_abs = (
-            hipify_result[s_abs].hipified_path
-            if (
-                s_abs in hipify_result
-                and hipify_result[s_abs].hipified_path is not None
-            )
-            else s_abs
-        )
+        hipified_s_abs = (hipify_result[s_abs].hipified_path if
+                          (s_abs in hipify_result
+                           and hipify_result[s_abs].hipified_path is not None)
+                          else s_abs)
         hipified_sources.append(hipified_s_abs)
 
-    assert len(hipified_sources) == len(args.sources)
+    assert (len(hipified_sources) == len(args.sources))
 
     # Print hipified source files.
     print("\n".join(hipified_sources))

csrc/cache_kernels.cu

@@ -16,7 +16,7 @@
 
 #include <algorithm>
 #include <cassert>
-#include <cfloat>
+#include <cfloat>  // FLT_MIN
 
 #ifdef USE_ROCM
   #include <hip/hip_bf16.h>
@@ -479,7 +479,6 @@ __global__ void concat_and_cache_ds_mla_kernel(
 
   // Compute the scale for the tile
   float tile_scale = max_abs / 448.f;
-  tile_scale = fmaxf(tile_scale, FLT_MIN);
 
   // The first lane of each half-warp writes the scale to kv_cache
   if ((lane_idx == 0) || (lane_idx == 16)) {

csrc/cpu/dnnl_helper.cpp

@@ -137,8 +137,9 @@ DNNLMatMulPrimitiveHandler::DNNLMatMulPrimitiveHandler(
 }
 
 void DNNLMatMulPrimitiveHandler::prepack_weight(
-    void* original_b_ptr, dnnl::memory::desc original_b_md,
-    dnnl::memory::desc b_target_mem_desc) {
+    void* original_b_ptr, dnnl::memory::desc b_target_mem_desc) {
+  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
+                                   {b_k_stride_, b_n_stride_});
   dnnl::memory original_weight(original_b_md, default_engine(), original_b_ptr);
   dnnl::memory packed_weight(b_target_mem_desc, default_engine());
   {
@@ -249,9 +250,7 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
   if (a_qs_ == QuantizationStrategy::PER_TOKEN) {
     assert(!use_azp_);
   };
-  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
-                                   {b_k_stride_, b_n_stride_});
-  prepack_weight(args.b_ptr, original_b_md,
+  prepack_weight(args.b_ptr,
                  create_primitive_desc(
                      MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
                                    .use_bias = false,
@@ -413,25 +412,12 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
   assert(ab_type_ == dnnl::memory::data_type::f32 ||
          ab_type_ == dnnl::memory::data_type::bf16 ||
          ab_type_ == dnnl::memory::data_type::f16);
-
-  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
-                                   {b_k_stride_, b_n_stride_});
-
-  prepack_weight(args.b_ptr, original_b_md,
+  prepack_weight(args.b_ptr,
                  create_primitive_desc(
-                     MSizeCacheKey{
-#ifdef VLLM_USE_ACL
-                         // Arm Compute Library (ACL) backend for oneDNN does
-                         // not support runtime
-                         // dimensions, so we set M to a default value
-                         .a_m_size = 128,
-                         .a_m_stride = b_k_size_,
-#else
-                         .a_m_size = DNNL_RUNTIME_DIM_VAL,
-                         .a_m_stride = DNNL_RUNTIME_DIM_VAL,
-#endif
-                         .use_bias = false,
-                         .bias_type = dnnl::memory::data_type::undef},
+                     MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
+                                   .a_m_stride = DNNL_RUNTIME_DIM_VAL,
+                                   .use_bias = false,
+                                   .bias_type = dnnl::memory::data_type::undef},
                      true)
                      .weights_desc());
   init_runtime_memory_cache(args);
@@ -457,30 +443,12 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) {
   c_storage->set_data_handle((void*)args.c_ptr);
   c_mem_desc->dims[0] = args.a_m_size;
 
-#ifndef VLLM_USE_ACL
-  // We do not support in ACL backend of oneDNN, we handle bias by:
-  // 1. copying it into the result tensor
-  // 2. attaching a fused-sum post-op to the matmul primitive
   if (args.use_bias) {
     auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(2);
     bias_storage->set_data_handle((void*)args.bias_ptr);
   }
-#endif
-  dnnl::matmul matmul = get_matmul_cache(args);
 
-// With ACL backend of oneDNN, the required memory format might change when the
-// source tensor dims change. This does not really happen in practice, so isn't
-// a performance hit, but we need to support it because the API allows for it.
-#ifdef VLLM_USE_ACL
-  auto new_expected_wei_desc =
-      dnnl::matmul::primitive_desc(
-          const_cast<dnnl_primitive_desc_t>(matmul.get_primitive_desc()))
-          .weights_desc();
-  if (new_expected_wei_desc != b_target_mem_desc_) {
-    prepack_weight(memory_cache_[DNNL_ARG_WEIGHTS].get_data_handle(),
-                   b_target_mem_desc_, new_expected_wei_desc);
-  }
-#endif
+  dnnl::matmul matmul = get_matmul_cache(args);
 
   auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(3);
   scratchpad_storage->set_data_handle(
@@ -516,13 +484,7 @@ dnnl::matmul::primitive_desc MatMulPrimitiveHandler::create_primitive_desc(
   } else {
     a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_,
                               {key.a_m_stride, 1});
-#ifdef VLLM_USE_ACL
-    // ACL's backend of oneDNN always expects the weight format to be "any"
-    b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_,
-                              dnnl::memory::format_tag::any);
-#else
     b_md = b_target_mem_desc_;
-#endif
   }
   dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_,
                           dnnl::memory::format_tag::ab);
@@ -532,18 +494,8 @@ dnnl::matmul::primitive_desc MatMulPrimitiveHandler::create_primitive_desc(
 
   if (key.use_bias) {
     dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1});
-// Since ACL's matmuls don't support passing a bias_md, we apply the bias
-// through a fused-sum post-op
-#ifdef VLLM_USE_ACL
-    dnnl::post_ops post_ops;
-    post_ops.append_sum();
-    attr.set_post_ops(post_ops);
-    return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
-                                        attr);
-#else
     return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md,
                                         c_md, attr);
-#endif
   } else {
     return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
                                         attr);
@@ -559,23 +511,13 @@ void MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
                    default_engine(), nullptr);
   set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get());
 
-// ACL matmuls don't support bias_md, so we don't need these
-#ifndef VLLM_USE_ACL
   memory_cache_[DNNL_ARG_BIAS] =
       dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
                    default_engine(), nullptr);
   set_runtime_memory_ptr(2, memory_cache_[DNNL_ARG_BIAS].get());
-#endif
+
   memory_cache_[DNNL_ARG_SCRATCHPAD] =
       dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
                    default_engine(), nullptr);
   set_runtime_memory_ptr(3, memory_cache_[DNNL_ARG_SCRATCHPAD].get());
 }
-
-bool is_onednn_acl_supported() {
-#ifdef VLLM_USE_ACL
-  return true;
-#else
-  return false;
-#endif
-}

csrc/cpu/dnnl_helper.h

@@ -101,7 +101,7 @@ class DNNLMatMulPrimitiveHandler {
  protected:
   DNNLMatMulPrimitiveHandler(const Args& args, dnnl::memory::data_type b_type);
 
-  void prepack_weight(void* original_b_ptr, dnnl::memory::desc original_b_md,
+  void prepack_weight(void* original_b_ptr,
                       dnnl::memory::desc b_target_mem_desc);
 
   void set_runtime_memory_ptr(size_t index, dnnl_memory* memory_ptr);

csrc/cpu/dnnl_kernels.cpp

@@ -527,42 +527,21 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
   MatMulPrimitiveHandler* ptr =
       reinterpret_cast<MatMulPrimitiveHandler*>(handler);
 
-// ACL matmuls expect contiguous source tensors
-#ifdef VLLM_USE_ACL
-  torch::Tensor a_contig = a.contiguous();
-#endif
-
   MatMulPrimitiveHandler::ExecArgs exec_args;
-
-#ifdef VLLM_USE_ACL
-  exec_args.a_m_size = a_contig.size(0);
-  exec_args.a_m_stride = a_contig.stride(0);
-#else
   exec_args.a_m_size = a.size(0);
   exec_args.a_m_stride = a.stride(0);
-#endif
+
   VLLM_DISPATCH_FLOATING_TYPES(a.scalar_type(), "onednn_mm", [&] {
     if (bias.has_value()) {
       exec_args.use_bias = true;
       exec_args.bias_type = get_dnnl_type<scalar_t>();
-#ifdef VLLM_USE_ACL
-      // ACL matmuls in oneDNN do not support a bias.
-      // We handle a matmul with bias by doing: c = bias; c += matmul(a, b)
-      c.copy_(bias.value());
-#else
       exec_args.bias_ptr = bias->data_ptr<scalar_t>();
-#endif
     } else {
       exec_args.use_bias = false;
       exec_args.bias_type = get_dnnl_type<void>();
       exec_args.bias_ptr = nullptr;
     }
-#ifdef VLLM_USE_ACL
-    exec_args.a_ptr = a_contig.data_ptr<scalar_t>();
-#else
     exec_args.a_ptr = a.data_ptr<scalar_t>();
-
-#endif
     exec_args.c_ptr = c.data_ptr<scalar_t>();
 
     ptr->execute(exec_args);

csrc/cpu/torch_bindings.cpp

@@ -27,8 +27,6 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b,
 void onednn_mm(torch::Tensor& c, const torch::Tensor& a,
                const std::optional<torch::Tensor>& bias, int64_t handler);
 
-bool is_onednn_acl_supported();
-
 void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
                         torch::Tensor& kv_cache, double scale,
                         torch::Tensor& block_tables, torch::Tensor& seq_lens);
@@ -183,9 +181,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int handler) -> ()");
   ops.impl("onednn_mm", torch::kCPU, &onednn_mm);
 
-  // Check if oneDNN was built with ACL backend
-  ops.def("is_onednn_acl_supported() -> bool", &is_onednn_acl_supported);
-
   // Create oneDNN W8A8 handler
   ops.def(
       "create_onednn_scaled_mm_handler(Tensor b, Tensor b_scales, ScalarType "

csrc/cutlass_extensions/vllm_cutlass_library_extension.py

@@ -27,7 +27,7 @@ VLLMDataTypeNames: dict[Union[VLLMDataType, DataType], str] = {
     **{
         VLLMDataType.u4b8: "u4b8",
         VLLMDataType.u8b128: "u8b128",
-    },
+    }
 }
 
 VLLMDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {
@@ -35,7 +35,7 @@ VLLMDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {
     **{
         VLLMDataType.u4b8: "cutlass::vllm_uint4b8_t",
         VLLMDataType.u8b128: "cutlass::vllm_uint8b128_t",
-    },
+    }
 }
 
 VLLMDataTypeSize: dict[Union[VLLMDataType, DataType], int] = {
@@ -43,7 +43,7 @@ VLLMDataTypeSize: dict[Union[VLLMDataType, DataType], int] = {
     **{
         VLLMDataType.u4b8: 4,
         VLLMDataType.u8b128: 8,
-    },
+    }
 }
 
 VLLMDataTypeVLLMScalarTypeTag: dict[Union[VLLMDataType, DataType], str] = {
@@ -67,13 +67,15 @@ VLLMDataTypeTorchDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {
     DataType.f32: "at::ScalarType::Float",
 }
 
-VLLMKernelScheduleTag: dict[
-    Union[MixedInputKernelScheduleType, KernelScheduleType], str
-] = {
-    **KernelScheduleTag,  # type: ignore
-    **{
-        MixedInputKernelScheduleType.TmaWarpSpecialized: "cutlass::gemm::KernelTmaWarpSpecialized",  # noqa: E501
-        MixedInputKernelScheduleType.TmaWarpSpecializedPingpong: "cutlass::gemm::KernelTmaWarpSpecializedPingpong",  # noqa: E501
-        MixedInputKernelScheduleType.TmaWarpSpecializedCooperative: "cutlass::gemm::KernelTmaWarpSpecializedCooperative",  # noqa: E501
-    },
-}
+VLLMKernelScheduleTag: dict[Union[
+    MixedInputKernelScheduleType, KernelScheduleType], str] = {
+        **KernelScheduleTag,  # type: ignore
+        **{
+            MixedInputKernelScheduleType.TmaWarpSpecialized:
+            "cutlass::gemm::KernelTmaWarpSpecialized",
+            MixedInputKernelScheduleType.TmaWarpSpecializedPingpong:
+            "cutlass::gemm::KernelTmaWarpSpecializedPingpong",
+            MixedInputKernelScheduleType.TmaWarpSpecializedCooperative:
+            "cutlass::gemm::KernelTmaWarpSpecializedCooperative",
+        }
+    }

csrc/moe/marlin_moe_wna16/generate_kernels.py

@@ -17,30 +17,25 @@ FILE_HEAD = """
 namespace MARLIN_NAMESPACE_NAME {
 """.strip()
 
-TEMPLATE = (
-    "template __global__ void Marlin<"
-    "{{scalar_t}}, "
-    "{{w_type_id}}, "
-    "{{s_type_id}}, "
-    "{{threads}}, "
-    "{{thread_m_blocks}}, "
-    "{{thread_n_blocks}}, "
-    "{{thread_k_blocks}}, "
-    "{{'true' if m_block_size_8 else 'false'}}, "
-    "{{stages}}, "
-    "{{group_blocks}}, "
-    "{{'true' if is_zp_float else 'false'}}>"
-    "( MARLIN_KERNEL_PARAMS );"
-)
+TEMPLATE = ("template __global__ void Marlin<"
+            "{{scalar_t}}, "
+            "{{w_type_id}}, "
+            "{{s_type_id}}, "
+            "{{threads}}, "
+            "{{thread_m_blocks}}, "
+            "{{thread_n_blocks}}, "
+            "{{thread_k_blocks}}, "
+            "{{'true' if m_block_size_8 else 'false'}}, "
+            "{{stages}}, "
+            "{{group_blocks}}, "
+            "{{'true' if is_zp_float else 'false'}}>"
+            "( MARLIN_KERNEL_PARAMS );")
 
 # int8 with zero point case (vllm::kU8) is also supported,
 # we don't add it to reduce wheel size.
 SCALAR_TYPES = [
-    "vllm::kU4",
-    "vllm::kU4B8",
-    "vllm::kU8B128",
-    "vllm::kFE4M3fn",
-    "vllm::kFE2M1f",
+    "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn",
+    "vllm::kFE2M1f"
 ]
 THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
 
@@ -63,12 +58,11 @@ def generate_new_kernels():
         all_template_str_list = []
 
         for group_blocks, m_blocks, thread_configs in itertools.product(
-            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
-        ):
+                GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):
+
             # act order case only support gptq-int4 and gptq-int8
             if group_blocks == 0 and scalar_type not in [
-                "vllm::kU4B8",
-                "vllm::kU8B128",
+                    "vllm::kU4B8", "vllm::kU8B128"
             ]:
                 continue
             if thread_configs[2] == 256:

csrc/quantization/gptq_marlin/generate_kernels.py

@@ -17,32 +17,28 @@ FILE_HEAD = """
 namespace MARLIN_NAMESPACE_NAME {
 """.strip()
 
-TEMPLATE = (
-    "template __global__ void Marlin<"
-    "{{scalar_t}}, "
-    "{{w_type_id}}, "
-    "{{s_type_id}}, "
-    "{{threads}}, "
-    "{{thread_m_blocks}}, "
-    "{{thread_n_blocks}}, "
-    "{{thread_k_blocks}}, "
-    "{{'true' if m_block_size_8 else 'false'}}, "
-    "{{stages}}, "
-    "{{group_blocks}}, "
-    "{{'true' if is_zp_float else 'false'}}>"
-    "( MARLIN_KERNEL_PARAMS );"
-)
+TEMPLATE = ("template __global__ void Marlin<"
+            "{{scalar_t}}, "
+            "{{w_type_id}}, "
+            "{{s_type_id}}, "
+            "{{threads}}, "
+            "{{thread_m_blocks}}, "
+            "{{thread_n_blocks}}, "
+            "{{thread_k_blocks}}, "
+            "{{'true' if m_block_size_8 else 'false'}}, "
+            "{{stages}}, "
+            "{{group_blocks}}, "
+            "{{'true' if is_zp_float else 'false'}}>"
+            "( MARLIN_KERNEL_PARAMS );")
 
 # int8 with zero point case (vllm::kU8) is also supported,
 # we don't add it to reduce wheel size.
 SCALAR_TYPES = [
-    "vllm::kU4",
-    "vllm::kU4B8",
-    "vllm::kU8B128",
-    "vllm::kFE4M3fn",
-    "vllm::kFE2M1f",
+    "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn",
+    "vllm::kFE2M1f"
 ]
-THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
+THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128),
+                  (128, 64, 128)]
 
 THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
 # group_blocks:
@@ -63,12 +59,11 @@ def generate_new_kernels():
         all_template_str_list = []
 
         for group_blocks, m_blocks, thread_configs in itertools.product(
-            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
-        ):
+                GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):
+
             # act order case only support gptq-int4 and gptq-int8
             if group_blocks == 0 and scalar_type not in [
-                "vllm::kU4B8",
-                "vllm::kU8B128",
+                    "vllm::kU4B8", "vllm::kU8B128"
             ]:
                 continue
             if thread_configs[2] == 256:
@@ -98,7 +93,8 @@ def generate_new_kernels():
             c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
 
             is_zp_float_list = [False]
-            if dtype == "fp16" and scalar_type == "vllm::kU4" and group_blocks == 4:
+            if dtype == "fp16" and scalar_type == "vllm::kU4" and \
+                    group_blocks == 4:
                 # HQQ (is_zp_float = true) only supports
                 # 4bit quantization and fp16
                 is_zp_float_list.append(True)

csrc/quantization/machete/generate.py

@@ -12,21 +12,20 @@ from functools import reduce
 from typing import Optional, Union
 
 import jinja2
-from vllm_cutlass_library_extension import (
-    DataType,
-    EpilogueScheduleTag,
-    EpilogueScheduleType,
-    MixedInputKernelScheduleType,
-    TileSchedulerTag,
-    TileSchedulerType,
-    VLLMDataType,
-    VLLMDataTypeNames,
-    VLLMDataTypeSize,
-    VLLMDataTypeTag,
-    VLLMDataTypeTorchDataTypeTag,
-    VLLMDataTypeVLLMScalarTypeTag,
-    VLLMKernelScheduleTag,
-)
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm_cutlass_library_extension import (DataType, EpilogueScheduleTag,
+                                            EpilogueScheduleType,
+                                            MixedInputKernelScheduleType,
+                                            TileSchedulerTag,
+                                            TileSchedulerType, VLLMDataType,
+                                            VLLMDataTypeNames,
+                                            VLLMDataTypeSize, VLLMDataTypeTag,
+                                            VLLMDataTypeTorchDataTypeTag,
+                                            VLLMDataTypeVLLMScalarTypeTag,
+                                            VLLMKernelScheduleTag)
+
+# yapf: enable
 
 #
 #   Generator templating
@@ -287,23 +286,18 @@ def generate_sch_sig(schedule_config: ScheduleConfig) -> str:
     tile_shape = (
         f"{schedule_config.tile_shape_mn[0]}x{schedule_config.tile_shape_mn[1]}"
     )
-    cluster_shape = (
-        f"{schedule_config.cluster_shape_mnk[0]}"
-        + f"x{schedule_config.cluster_shape_mnk[1]}"
-        + f"x{schedule_config.cluster_shape_mnk[2]}"
-    )
-    kernel_schedule = VLLMKernelScheduleTag[schedule_config.kernel_schedule].split(
-        "::"
-    )[-1]
-    epilogue_schedule = EpilogueScheduleTag[schedule_config.epilogue_schedule].split(
-        "::"
-    )[-1]
-    tile_scheduler = TileSchedulerTag[schedule_config.tile_scheduler].split("::")[-1]
+    cluster_shape = (f"{schedule_config.cluster_shape_mnk[0]}" +
+                     f"x{schedule_config.cluster_shape_mnk[1]}" +
+                     f"x{schedule_config.cluster_shape_mnk[2]}")
+    kernel_schedule = VLLMKernelScheduleTag[schedule_config.kernel_schedule]\
+        .split("::")[-1]
+    epilogue_schedule = EpilogueScheduleTag[
+        schedule_config.epilogue_schedule].split("::")[-1]
+    tile_scheduler = TileSchedulerTag[schedule_config.tile_scheduler]\
+        .split("::")[-1]
 
-    return (
-        f"{tile_shape}_{cluster_shape}_{kernel_schedule}"
-        + f"_{epilogue_schedule}_{tile_scheduler}"
-    )
+    return (f"{tile_shape}_{cluster_shape}_{kernel_schedule}" +
+            f"_{epilogue_schedule}_{tile_scheduler}")
 
 
 # mostly unique shorter sch_sig
@@ -322,24 +316,18 @@ def generate_terse_sch_sig(schedule_config: ScheduleConfig) -> str:
 
 # unique type_name
 def generate_type_signature(kernel_types: TypeConfig):
-    return str(
-        "".join(
-            [
-                VLLMDataTypeNames[getattr(kernel_types, field.name)]
-                for field in fields(TypeConfig)
-            ]
-        )
-    )
+    return str("".join([
+        VLLMDataTypeNames[getattr(kernel_types, field.name)]
+        for field in fields(TypeConfig)
+    ]))
 
 
 def generate_type_option_name(kernel_types: TypeConfig):
-    return ", ".join(
-        [
-            f"{field.name.replace('b_', 'with_') + '_type'}="
-            + VLLMDataTypeNames[getattr(kernel_types, field.name)]
-            for field in fields(TypeConfig)
-        ]
-    )
+    return ", ".join([
+        f"{field.name.replace('b_', 'with_')+'_type'}=" +
+        VLLMDataTypeNames[getattr(kernel_types, field.name)]
+        for field in fields(TypeConfig)
+    ])
 
 
 def is_power_of_two(n):
@@ -347,6 +335,7 @@ def is_power_of_two(n):
 
 
 def to_cute_constant(value: list[int]):
+
     def _to_cute_constant(value: int):
         if is_power_of_two(value):
             return f"_{value}"
@@ -361,11 +350,11 @@ def to_cute_constant(value: list[int]):
 
 def unique_schedules(impl_configs: list[ImplConfig]):
     # Use dict over set for deterministic ordering
-    return list(
-        {
-            sch: None for impl_config in impl_configs for sch in impl_config.schedules
-        }.keys()
-    )
+    return list({
+        sch: None
+        for impl_config in impl_configs
+        for sch in impl_config.schedules
+    }.keys())
 
 
 def unsigned_type_with_bitwidth(num_bits):
@@ -391,7 +380,7 @@ template_globals = {
     "gen_type_sig": generate_type_signature,
     "unique_schedules": unique_schedules,
     "unsigned_type_with_bitwidth": unsigned_type_with_bitwidth,
-    "gen_type_option_name": generate_type_option_name,
+    "gen_type_option_name": generate_type_option_name
 }
 
 
@@ -409,28 +398,23 @@ prepack_dispatch_template = create_template(PREPACK_TEMPLATE)
 def create_sources(impl_configs: list[ImplConfig], num_impl_files=8):
     sources = []
 
-    sources.append(
-        (
-            "machete_mm_dispatch",
-            mm_dispatch_template.render(impl_configs=impl_configs),
-        )
-    )
+    sources.append((
+        "machete_mm_dispatch",
+        mm_dispatch_template.render(impl_configs=impl_configs),
+    ))
 
     prepack_types = []
     for impl_config in impl_configs:
-        convert_type = (
-            impl_config.types.a
-            if impl_config.types.b_group_scale == DataType.void
-            else impl_config.types.b_group_scale
-        )
+        convert_type = impl_config.types.a \
+             if impl_config.types.b_group_scale == DataType.void \
+             else impl_config.types.b_group_scale
         prepack_types.append(
             PrepackTypeConfig(
                 a=impl_config.types.a,
                 b_num_bits=VLLMDataTypeSize[impl_config.types.b],
                 convert=convert_type,
                 accumulator=impl_config.types.accumulator,
-            )
-        )
+            ))
 
     def prepacked_type_key(prepack_type: PrepackTypeConfig):
         # For now, we can just use the first accumulator type seen since
@@ -446,14 +430,10 @@ def create_sources(impl_configs: list[ImplConfig], num_impl_files=8):
             unique_prepack_types.append(prepack_type)
             prepack_types_seen.add(key)
 
-    sources.append(
-        (
-            "machete_prepack",
-            prepack_dispatch_template.render(
-                types=unique_prepack_types,
-            ),
-        )
-    )
+    sources.append((
+        "machete_prepack",
+        prepack_dispatch_template.render(types=unique_prepack_types, ),
+    ))
 
     # Split up impls across files
     num_impls = reduce(lambda x, y: x + len(y.schedules), impl_configs, 0)
@@ -486,12 +466,10 @@ def create_sources(impl_configs: list[ImplConfig], num_impl_files=8):
         curr_impl_in_file += len(files_impls[-1][-1].schedules)
 
     for part, file_impls in enumerate(files_impls):
-        sources.append(
-            (
-                f"machete_mm_impl_part{part + 1}",
-                mm_impl_template.render(impl_configs=file_impls),
-            )
-        )
+        sources.append((
+            f"machete_mm_impl_part{part+1}",
+            mm_impl_template.render(impl_configs=file_impls),
+        ))
 
     return sources
 
@@ -536,7 +514,8 @@ def generate():
     # For now we use the same heuristic for all types
     # Heuristic is currently tuned for H100s
     default_heuristic = [
-        (cond, ScheduleConfig(*tile_config, **sch_common_params))  # type: ignore
+        (cond, ScheduleConfig(*tile_config,
+                              **sch_common_params))  # type: ignore
         for cond, tile_config in default_tile_heuristic_config.items()
     ]
 
@@ -562,18 +541,14 @@ def generate():
             a_token_scale=DataType.void,
             out=a,
             accumulator=DataType.f32,
-        )
-        for b in (VLLMDataType.u4b8, VLLMDataType.u8b128)
-        for a in (DataType.f16, DataType.bf16)
-    )
+        ) for b in (VLLMDataType.u4b8, VLLMDataType.u8b128)
+        for a in (DataType.f16, DataType.bf16))
 
     impl_configs += [
         ImplConfig(x[0], x[1], x[2])
-        for x in zip(
-            GPTQ_kernel_type_configs,
-            itertools.repeat(get_unique_schedules(default_heuristic)),
-            itertools.repeat(default_heuristic),
-        )
+        for x in zip(GPTQ_kernel_type_configs,
+                     itertools.repeat(get_unique_schedules(default_heuristic)),
+                     itertools.repeat(default_heuristic))
     ]
 
     AWQ_kernel_type_configs = list(
@@ -586,18 +561,14 @@ def generate():
             a_token_scale=DataType.void,
             out=a,
             accumulator=DataType.f32,
-        )
-        for b in (DataType.u4, DataType.u8)
-        for a in (DataType.f16, DataType.bf16)
-    )
+        ) for b in (DataType.u4, DataType.u8)
+        for a in (DataType.f16, DataType.bf16))
 
     impl_configs += [
         ImplConfig(x[0], x[1], x[2])
-        for x in zip(
-            AWQ_kernel_type_configs,
-            itertools.repeat(get_unique_schedules(default_heuristic)),
-            itertools.repeat(default_heuristic),
-        )
+        for x in zip(AWQ_kernel_type_configs,
+                     itertools.repeat(get_unique_schedules(default_heuristic)),
+                     itertools.repeat(default_heuristic))
     ]
 
     # TODO: Support W4A8 when ready

docs/configuration/conserving_memory.md

@@ -53,7 +53,7 @@ llm = LLM(model="adept/fuyu-8b",
 By default, we optimize model inference using CUDA graphs which take up extra memory in the GPU.
 
 !!! warning
-    CUDA graph capture increases GPU memory usage. Adjust capture sizes if you need to conserve memory.
+    CUDA graph capture takes up more memory in V1 than in V0.
 
 You can adjust `compilation_config` to achieve a better balance between inference speed and memory usage:
 

docs/configuration/optimization.md

@@ -33,7 +33,7 @@ In vLLM V1, the default preemption mode is `RECOMPUTE` rather than `SWAP`, as re
 
 Chunked prefill allows vLLM to process large prefills in smaller chunks and batch them together with decode requests. This feature helps improve both throughput and latency by better balancing compute-bound (prefill) and memory-bound (decode) operations.
 
-In vLLM V1, **chunked prefill is always enabled by default** so that behavior is consistent across supported models.
+In vLLM V1, **chunked prefill is always enabled by default**. This is different from vLLM V0, where it was conditionally enabled based on model characteristics.
 
 With chunked prefill enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill operations. When there are available tokens in the `max_num_batched_tokens` budget, it schedules pending prefills. If a pending prefill request cannot fit into `max_num_batched_tokens`, it automatically chunks it.
 
@@ -49,7 +49,7 @@ You can tune the performance by adjusting `max_num_batched_tokens`:
 - Smaller values (e.g., 2048) achieve better inter-token latency (ITL) because there are fewer prefills slowing down decodes.
 - Higher values achieve better time to first token (TTFT) as you can process more prefill tokens in a batch.
 - For optimal throughput, we recommend setting `max_num_batched_tokens > 8192` especially for smaller models on large GPUs.
-- If `max_num_batched_tokens` is the same as `max_model_len`, the scheduler behaves similarly to the legacy policy where large prefills ran without chunking (while still prioritizing decodes).
+- If `max_num_batched_tokens` is the same as `max_model_len`, that's almost the equivalent to the V0 default scheduling policy (except that it still prioritizes decodes).
 
 ```python
 from vllm import LLM

docs/contributing/model/basic.md

@@ -133,7 +133,8 @@ We consider 3 different scenarios:
 For case (1), we recommend looking at the implementation of [`MambaForCausalLM`](gh-file:vllm/model_executor/models/mamba.py) (for Mamba-1) or [`Mamba2ForCausalLM`](gh-file:vllm/model_executor/models/mamba2.py) (for Mamba-2) as a reference.
 The model should inherit protocol `IsAttentionFree` and also implement class methods `get_mamba_state_dtype_from_config` and `get_mamba_state_shape_from_config` to calculate the state shapes and data types from the config.
 For the mamba layers themselves, please use the [`MambaMixer`](gh-file:vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](gh-file:vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
-Please avoid reintroducing legacy cache managers such as `MambaCacheManager` or any previously removed code paths from older implementations.
+Please *do not* use the `MambaCacheManager` (deprecated in V1) or replicate any of the V0-specific code paths in the existing model implementations.
+V0-only classes and code will be removed in the very near future.
 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in <gh-file:vllm/model_executor/models/config.py> to ensure that the runtime defaults are optimized.
 
 For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](gh-file:vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](gh-file:vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).

docs/deployment/frameworks/hf_inference_endpoints.md

@@ -61,7 +61,7 @@ This is the easiest way to get started with vLLM on Hugging Face Inference Endpo
 
 ### Method 2: Guided Deployment (Transformers Models)
 
-This method applies to models with the [`transformers` library tag](https://huggingface.co/models?library=transformers) in their metadata. It allows you to deploy a model directly from the Hub UI without manual configuration.
+This method applies to models with the `transformers` library tag in their metadata. It allows you to deploy a model directly from the Hub UI without manual configuration.
 
 1. Navigate to a model on [Hugging Face Hub](https://huggingface.co/models).  
    For this example we will use the [`ibm-granite/granite-docling-258M`](https://huggingface.co/ibm-granite/granite-docling-258M) model. You can verify that the model is compatible by checking the front matter in the [README](https://huggingface.co/ibm-granite/granite-docling-258M/blob/main/README.md), where the library is tagged as `library: transformers`.
@@ -128,7 +128,7 @@ Some models require manual deployment because they:
 
 These models cannot be deployed using the **Deploy** button on the model card.
 
-In this guide, we demonstrate manual deployment using the [`rednote-hilab/dots.ocr`](https://huggingface.co/rednote-hilab/dots.ocr) model, an OCR model integrated with vLLM (see vLLM [PR](https://github.com/vllm-project/vllm/pull/24645)).
+In this guide, we demonstrate manual deployment using the [rednote-hilab/dots.ocr](https://huggingface.co/rednote-hilab/dots.ocr) model, an OCR model integrated with vLLM (see vLLM [PR](https://github.com/vllm-project/vllm/pull/24645)).
 
 1. Start a new deployment. Go to [Inference Endpoints](https://endpoints.huggingface.co/) and click `New`.
 

docs/deployment/integrations/kaito.md

@@ -1,5 +0,0 @@
-# KAITO
-
-[KAITO](https://kaito-project.github.io/kaito/docs/) is a Kubernetes operator that supports deploying and serving LLMs with vLLM. It offers managing large models via container images with built-in OpenAI-compatible inference, auto-provisioning GPU nodes and curated model presets.
-
-Please refer to [quick start](https://kaito-project.github.io/kaito/docs/quick-start) for more details.

docs/deployment/integrations/production-stack.md

@@ -55,7 +55,7 @@ sudo kubectl port-forward svc/vllm-router-service 30080:80
 And then you can send out a query to the OpenAI-compatible API to check the available models:
 
 ```bash
-curl -o- http://localhost:30080/v1/models
+curl -o- http://localhost:30080/models
 ```
 
 ??? console "Output"
@@ -78,7 +78,7 @@ curl -o- http://localhost:30080/v1/models
 To send an actual chatting request, you can issue a curl request to the OpenAI `/completion` endpoint:
 
 ```bash
-curl -X POST http://localhost:30080/v1/completions \
+curl -X POST http://localhost:30080/completions \
   -H "Content-Type: application/json" \
   -d '{
     "model": "facebook/opt-125m",

docs/deployment/k8s.md

@@ -12,7 +12,6 @@ Alternatively, you can deploy vLLM to Kubernetes using any of the following:
 
 - [Helm](frameworks/helm.md)
 - [InftyAI/llmaz](integrations/llmaz.md)
-- [KAITO](integrations/kaito.md)
 - [KServe](integrations/kserve.md)
 - [KubeRay](integrations/kuberay.md)
 - [kubernetes-sigs/lws](frameworks/lws.md)

docs/design/metrics.md

@@ -1,12 +1,12 @@
 # Metrics
 
-vLLM exposes a rich set of metrics to support observability and capacity planning for the V1 engine.
+Ensure the v1 LLM Engine exposes a superset of the metrics available in v0.
 
 ## Objectives
 
-- Provide comprehensive coverage of engine and request level metrics to aid production monitoring.
-- Prioritize Prometheus integrations, as this is what we expect to be used in production environments.
-- Offer logging support (i.e. printing metrics to the info log) for ad-hoc testing, debugging, development, and exploratory use cases.
+- Achieve parity of metrics between v0 and v1.
+- The priority use case is accessing these metrics via Prometheus, as this is what we expect to be used in production environments.
+- Logging support (i.e. printing metrics to the info log) is provided for more ad-hoc testing, debugging, development, and exploratory use cases.
 
 ## Background
 
@@ -17,9 +17,9 @@ Metrics in vLLM can be categorized as follows:
 
 The mental model is that server-level metrics help explain the values of request-level metrics.
 
-### Metrics Overview
+### v0 Metrics
 
-The following metrics are exposed via a Prometheus-compatible `/metrics` endpoint using the `vllm:` prefix and are documented under [Inferencing and Serving -> Production Metrics](../usage/metrics.md):
+In v0, the following metrics are exposed via a Prometheus-compatible `/metrics` endpoint using the `vllm:` prefix:
 
 - `vllm:num_requests_running` (Gauge)
 - `vllm:num_requests_swapped` (Gauge)
@@ -57,6 +57,8 @@ The following metrics are exposed via a Prometheus-compatible `/metrics` endpoin
 - `vllm:spec_decode_num_draft_tokens_total` (Counter)
 - `vllm:spec_decode_num_emitted_tokens_total` (Counter)
 
+These are documented under [Inferencing and Serving -> Production Metrics](../usage/metrics.md).
+
 ### Grafana Dashboard
 
 vLLM also provides [a reference example](../examples/online_serving/prometheus_grafana.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
@@ -84,7 +86,7 @@ See [the PR which added this Dashboard](gh-pr:2316) for interesting and useful b
 
 Prometheus support was initially added [using the aioprometheus library](gh-pr:1890), but a switch was made quickly to [prometheus_client](gh-pr:2730). The rationale is discussed in both linked PRs.
 
-During those migrations we briefly lost a `MetricsMiddleware` to track HTTP metrics, but this was reinstated [using prometheus_fastapi_instrumentator](gh-pr:15657):
+With the switch to `aioprometheus`, we lost a `MetricsMiddleware` to track HTTP metrics, but this was reinstated [using prometheus_fastapi_instrumentator](gh-pr:15657):
 
 ```bash
 $ curl http://0.0.0.0:8000/metrics 2>/dev/null  | grep -P '^http_(?!.*(_bucket|_created|_sum)).*'
@@ -95,6 +97,10 @@ http_request_duration_highr_seconds_count 201.0
 http_request_duration_seconds_count{handler="/v1/completions",method="POST"} 201.0
 ```
 
+### Multi-process Mode
+
+In v0, metrics are collected in the engine core process and we use multiprocess mode to make them available in the API server process. See <gh-pr:7279>.
+
 ### Built in Python/Process Metrics
 
 The following metrics are supported by default by `prometheus_client`, but they are not exposed when multiprocess mode is used:
@@ -110,7 +116,22 @@ The following metrics are supported by default by `prometheus_client`, but they
 - `process_open_fds`
 - `process_max_fds`
 
-This is relevant because if we move away from multiprocess mode we get these back. However, it's questionable how relevant these are if they don't aggregate these stats for all processes that make up a vLLM instance.
+This is relevant because if we move away from multiprocess mode in v1,
+we get these back. However, it's questionable how relevant these are
+if they don't aggregate these stats for all processes that make up a
+vLLM instance.
+
+### v0 PRs and Issues
+
+For background, these are some of the relevant PRs which added the v0 metrics:
+
+- <gh-pr:1890>
+- <gh-pr:2316>
+- <gh-pr:2730>
+- <gh-pr:4464>
+- <gh-pr:7279>
+
+Also note the ["Even Better Observability"](gh-issue:3616) feature where e.g. [a detailed roadmap was laid out](gh-issue:3616#issuecomment-2030858781).
 
 ## v1 Design
 
@@ -375,8 +396,9 @@ recent metric is used, but only from currently running processes.
 
 This was added in <gh-pr:9477> and there is
 [at least one known user](https://github.com/kubernetes-sigs/gateway-api-inference-extension/pull/54).
-If we revisit this design and deprecate the old metric, we should
-coordinate with downstream users so they can migrate before the removal.
+If we revisit this design and deprecate the old metric, we should reduce
+the need for a significant deprecation period by making the change in
+v0 also and asking this project to move to the new metric.
 
 ### Prefix Cache metrics
 
@@ -469,7 +491,7 @@ if seq_group.is_finished():
 
 This seems duplicative, and one of them should be removed. The latter
 is used by the Grafana dashboard, so we should deprecate or remove the
-former.
+former from v0.
 
 ### Prefix Cache Hit Rate
 
@@ -478,7 +500,7 @@ See above - we now expose 'queries' and 'hits' counters rather than a
 
 ### KV Cache Offloading
 
-Two legacy metrics relate to a "swapped" preemption mode that is no
+Two v0 metrics relate to a "swapped" preemption mode that is no
 longer relevant in v1:
 
 - `vllm:num_requests_swapped`
@@ -489,7 +511,7 @@ cache to complete other requests), we swap kv cache blocks out to CPU
 memory. This is also known as "KV cache offloading" and is configured
 with `--swap-space` and `--preemption-mode`.
 
-Historically, [vLLM has long supported beam search](gh-issue:6226). The
+In v0, [vLLM has long supported beam search](gh-issue:6226). The
 SequenceGroup encapsulated the idea of N Sequences which
 all shared the same prompt kv blocks. This enabled KV cache block
 sharing between requests, and copy-on-write to do branching. CPU
@@ -502,7 +524,7 @@ and the part of the prompt that was evicted can be recomputed.
 
 SequenceGroup was removed in V1, although a replacement will be
 required for "parallel sampling" (`n>1`).
-[Beam search was moved out of the core](gh-issue:8306). There was a
+[Beam search was moved out of the core (in V0)](gh-issue:8306). There was a
 lot of complex code for a very uncommon feature.
 
 In V1, with prefix caching being better (zero over head) and therefore
@@ -513,7 +535,7 @@ better.
 
 ### Parallel Sampling
 
-Some legacy metrics are only relevant in the context of "parallel
+Some v0 metrics are only relevant in the context of "parallel
 sampling". This is where the `n` parameter in a request is used to
 request multiple completions from the same prompt.
 
@@ -532,7 +554,7 @@ also add these metrics.
 
 ### Speculative Decoding
 
-Some legacy metrics are specific to "speculative decoding". This is where
+Some v0 metrics are specific to "speculative decoding". This is where
 we generate candidate tokens using a faster, approximate method or
 model and then validate those tokens with the larger model.
 
@@ -544,7 +566,7 @@ model and then validate those tokens with the larger model.
 
 There is a PR under review (<gh-pr:12193>) to add "prompt lookup (ngram)"
 speculative decoding to v1. Other techniques will follow. We should
-revisit these metrics in this context.
+revisit the v0 metrics in this context.
 
 !!! note
     We should probably expose acceptance rate as separate accepted
@@ -617,7 +639,7 @@ metrics are often relatively straightforward to add:
    metrics are usually of very limited use unless they can be enabled
    by default and in production.
 3. They have an impact on development and maintenance of the
-   project. Every metric added over time has made this effort more
+   project. Every metric added to v0 has made this v1 effort more
    time-consuming, and perhaps not all metrics justify this ongoing
    investment in their maintenance.
 
@@ -628,7 +650,7 @@ performance and health. Tracing, on the other hand, tracks individual
 requests as they move through different services and components. Both
 fall under the more general heading of "Observability".
 
-vLLM has support for OpenTelemetry tracing:
+v0 has support for OpenTelemetry tracing:
 
 - Added by <gh-pr:4687>
 - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
@@ -641,11 +663,11 @@ OpenTelemetry has a
 [Gen AI Working Group](https://github.com/open-telemetry/community/blob/main/projects/gen-ai.md).
 
 Since metrics is a big enough topic on its own, we are going to tackle
-the topic of tracing separately.
+the topic of tracing in v1 separately.
 
 ### OpenTelemetry Model Forward vs Execute Time
 
-The current implementation exposes the following two metrics:
+In v0, we have the following two metrics:
 
 - `vllm:model_forward_time_milliseconds` (Histogram) - The time spent
   in the model forward pass when this request was in the batch.

docs/design/moe_kernel_features.md

@@ -93,8 +93,6 @@ To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels
 | gpt oss triton               | standard              | N/A              | N/A           | <sup>5</sup>                                                | Y                     | Y       | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts]                                                                    |
 | deep gemm+triton<sup>2</sup> | standard,</br>batched | all<sup>1</sup>  | G(128),A,T    | silu, gelu                                                  | <sup>6</sup>          | Y       | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],</br>[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts]                                                 |
 | marlin                       | standard              | <sup>3</sup>     | <sup>3</sup>  | silu,</br>swigluoai                                         | Y                     | N       | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe]                                                                                                                                                                                                                |
-
-| marlin experts               | standard              | N/A              | N/A           | silu,</br>swigluoai                                         | Y                     | Y       | [`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts]                                                                                                                                                                                                                      |
 | trtllm                       | standard              | mxfp4,</br>nvfp4 | G(16),G(32)   | <sup>5</sup>                                                | N                     | Y       | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts]                                                                                                                                                                                                                      |
 | pallas                       | standard              | N/A              | N/A           | silu                                                        | N                     | N       | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe]                                                                                                                                                                                                                                    |
 | iterative                    | standard              | N/A              | N/A           | silu                                                        | N                     | N       | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe]                                                                                                                                                                                                                           |
@@ -116,6 +114,6 @@ The following table shows "families" of modular kernels that are intended to wor
 
 | backend                          | `FusedMoEPrepareAndFinalize` subclasses                    | `FusedMoEPermuteExpertsUnpermute` subclasses                                                                               |
 |----------------------------------|------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------|
-| deepep_high_throughput           | `DeepEPHTPrepareAndFinalize`                               |  `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts`      |
-| deepep_low_latency,</br>pplx     | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` |  `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8`|
+| deepep_high_throughput,</br>pplx | `DeepEPHTPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8` |
+| deepep_low_latency               | `DeepEPLLPrepareAndFinalize`                               | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`                             |
 | flashinfer                       | `FlashInferCutlassMoEPrepareAndFinalize`                   | `FlashInferExperts`                                                                                                        |

docs/design/multiprocessing.md

@@ -60,6 +60,30 @@ Multiple vLLM dependencies indicate either a preference or requirement for using
 It is perhaps more accurate to say that there are known problems with using
 `fork` after initializing these dependencies.
 
+## Current State (v0)
+
+The environment variable `VLLM_WORKER_MULTIPROC_METHOD` can be used to control which method is used by vLLM. The current default is `fork`.
+
+- <https://github.com/vllm-project/vllm/blob/d05f88679bedd73939251a17c3d785a354b2946c/vllm/envs.py#L339-L342>
+
+When we know we own the process because the `vllm` command was used, we use
+`spawn` because it's the most widely compatible.
+
+- <https://github.com/vllm-project/vllm/blob/d05f88679bedd73939251a17c3d785a354b2946c/vllm/scripts.py#L123-L140>
+
+The `multiproc_xpu_executor` forces the use of `spawn`.
+
+- <https://github.com/vllm-project/vllm/blob/d05f88679bedd73939251a17c3d785a354b2946c/vllm/executor/multiproc_xpu_executor.py#L14-L18>
+
+There are other miscellaneous places hard-coding the use of `spawn`:
+
+- <https://github.com/vllm-project/vllm/blob/d05f88679bedd73939251a17c3d785a354b2946c/vllm/distributed/device_communicators/all_reduce_utils.py#L135>
+- <https://github.com/vllm-project/vllm/blob/d05f88679bedd73939251a17c3d785a354b2946c/vllm/entrypoints/openai/api_server.py#L184>
+
+Related PRs:
+
+- <gh-pr:8823>
+
 ## Prior State in v1
 
 There was an environment variable to control whether multiprocessing is used in

docs/design/plugin_system.md

@@ -49,7 +49,7 @@ Every plugin has three parts:
 
 - **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported.
 
-- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre/post processing of the model prompt and model output for pooling models. The plugin function returns the IOProcessor's class fully qualified name.
+- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre/post processing of the model prompt and model output for poling models. The plugin function returns the IOProcessor's class fully qualified name.
 
 ## Guidelines for Writing Plugins
 

docs/design/prefix_caching.md

@@ -94,6 +94,9 @@ To improve privacy in shared environments, vLLM supports isolating prefix cache
 
 With this setup, cache sharing is limited to users or requests that explicitly agree on a common salt, enabling cache reuse within a trust group while isolating others.
 
+!!! note
+    Cache isolation is not supported in engine V0.
+
 ## Data Structure
 
 The prefix caching in vLLM v1 is implemented in the KV cache manager. The basic building block is the “Block” data class (simplified):
@@ -186,7 +189,7 @@ Time 1:
   Cache Blocks: 0, 1, 3
 ```
 
-As can be seen, block 3 is a new full block and is cached. However, it is redundant as block 1, meaning that we cached the same block twice. Because the block table in vLLM v1 is append-only, changing the block table from `[0, 3]` to `[0, 1]` is not allowed. As a result, we will have duplicated blocks for the hash key E-H. This duplication will be eliminated when the request is freed.
+As can be seen, block 3 is a new full block and is cached. However, it is redundant as block 1, meaning that we cached the same block twice. In v0, when detecting block 3 is duplicated, we free block 3 and let Request 2 use block 1 instead, so its block table becomes `[0, 1]` in Time 1. However, the block table in vLLM v1 is append-only, meaning that changing the block table from `[0, 3]` to `[0, 1]` is not allowed. As a result, we will have duplicated blocks for the hash key E-H. This duplication will be eliminated when the request is freed.
 
 ### Free
 

docs/features/custom_logitsprocs.md

@@ -166,7 +166,7 @@ The `DummyLogitsProcessor.update_state()` implementation maintains a "sparse" re
 
 ### Wrapping an Existing Request-Level Logits Processor
 
-Although the vLLM engine applies logits processors at batch granularity, some users may want to use vLLM with a "request-level" logits processor implementation - an implementation which operates on individual requests. Earlier request-level processors were implemented as `Callable` objects conforming to the following type annotation:
+Although the vLLM engine applies logits processors at batch granularity, some users may want to use vLLM with a "request-level" logits processor implementation - an implementation which operates on individual requests. This will be especially true if your logits processor was developed for vLLM version 0, which required it to be a `Callable` (as described [here](https://docs.vllm.ai/en/v0.10.1.1/api/vllm/logits_process.html)) conforming to the following type annotation:
 
 ``` python
 RequestLogitsProcessor = Union[

docs/features/spec_decode.md

@@ -16,8 +16,8 @@ Speculative decoding is a technique which improves inter-token latency in memory
 The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
 
 !!! warning
-    Speculative decoding with a draft model requires the V1 engine.
-    Older releases that predate V1 (such as the 0.10.x series) raise a `NotImplementedError`.
+    In vllm v0.10.0, speculative decoding with a draft model is not supported.
+    If you use the following code, you will get a `NotImplementedError`.
 
 ??? code
 

docs/features/tool_calling.md

@@ -191,14 +191,10 @@ VLLM also provides a pythonic and JSON-based chat template for Llama 4, but pyth
 
 For Llama 4 model, use `--tool-call-parser llama4_pythonic --chat-template examples/tool_chat_template_llama4_pythonic.jinja`.
 
-### IBM Granite
+#### IBM Granite
 
 Supported models:
 
-* `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models
-
-    Recommended flags: `--tool-call-parser hermes`
-
 * `ibm-granite/granite-3.0-8b-instruct`
 
     Recommended flags: `--tool-call-parser granite --chat-template examples/tool_chat_template_granite.jinja`

docs/mkdocs/hooks/generate_argparse.py

@@ -33,11 +33,8 @@ def auto_mock(module, attr, max_mocks=50):
         try:
             # First treat attr as an attr, then as a submodule
             with patch("importlib.metadata.version", return_value="0.0.0"):
-                return getattr(
-                    importlib.import_module(module),
-                    attr,
-                    importlib.import_module(f"{module}.{attr}"),
-                )
+                return getattr(importlib.import_module(module), attr,
+                               importlib.import_module(f"{module}.{attr}"))
         except importlib.metadata.PackageNotFoundError as e:
             raise e
         except ModuleNotFoundError as e:
@@ -45,8 +42,7 @@ def auto_mock(module, attr, max_mocks=50):
             sys.modules[e.name] = PydanticMagicMock()
 
     raise ImportError(
-        f"Failed to import {module}.{attr} after mocking {max_mocks} imports"
-    )
+        f"Failed to import {module}.{attr} after mocking {max_mocks} imports")
 
 
 latency = auto_mock("vllm.benchmarks", "latency")
@@ -65,7 +61,9 @@ class MarkdownFormatter(HelpFormatter):
     """Custom formatter that generates markdown for argument groups."""
 
     def __init__(self, prog, starting_heading_level=3):
-        super().__init__(prog, max_help_position=float("inf"), width=float("inf"))
+        super().__init__(prog,
+                         max_help_position=float('inf'),
+                         width=float('inf'))
         self._section_heading_prefix = "#" * starting_heading_level
         self._argument_heading_prefix = "#" * (starting_heading_level + 1)
         self._markdown_output = []
@@ -87,19 +85,23 @@ class MarkdownFormatter(HelpFormatter):
 
     def add_arguments(self, actions):
         for action in actions:
-            if len(action.option_strings) == 0 or "--help" in action.option_strings:
+            if (len(action.option_strings) == 0
+                    or "--help" in action.option_strings):
                 continue
 
-            option_strings = f"`{'`, `'.join(action.option_strings)}`"
+            option_strings = f'`{"`, `".join(action.option_strings)}`'
             heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
             self._markdown_output.append(heading_md)
 
             if choices := action.choices:
-                choices = f"`{'`, `'.join(str(c) for c in choices)}`"
-                self._markdown_output.append(f"Possible choices: {choices}\n\n")
-            elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
-                metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
-                self._markdown_output.append(f"Possible choices: {metavar}\n\n")
+                choices = f'`{"`, `".join(str(c) for c in choices)}`'
+                self._markdown_output.append(
+                    f"Possible choices: {choices}\n\n")
+            elif ((metavar := action.metavar)
+                  and isinstance(metavar, (list, tuple))):
+                metavar = f'`{"`, `".join(str(m) for m in metavar)}`'
+                self._markdown_output.append(
+                    f"Possible choices: {metavar}\n\n")
 
             if action.help:
                 self._markdown_output.append(f"{action.help}\n\n")
@@ -114,7 +116,7 @@ class MarkdownFormatter(HelpFormatter):
 
 def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser:
     """Create a parser for the given class with markdown formatting.
-
+    
     Args:
         cls: The class to create a parser for
         **kwargs: Additional keyword arguments to pass to `cls.add_cli_args`.
@@ -141,17 +143,24 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
 
     # Create parsers to document
     parsers = {
-        "engine_args": create_parser(EngineArgs.add_cli_args),
-        "async_engine_args": create_parser(
-            AsyncEngineArgs.add_cli_args, async_args_only=True
-        ),
-        "serve": create_parser(cli_args.make_arg_parser),
-        "chat": create_parser(ChatCommand.add_cli_args),
-        "complete": create_parser(CompleteCommand.add_cli_args),
-        "bench_latency": create_parser(latency.add_cli_args),
-        "bench_throughput": create_parser(throughput.add_cli_args),
-        "bench_serve": create_parser(serve.add_cli_args),
-        "run-batch": create_parser(run_batch.make_arg_parser),
+        "engine_args":
+        create_parser(EngineArgs.add_cli_args),
+        "async_engine_args":
+        create_parser(AsyncEngineArgs.add_cli_args, async_args_only=True),
+        "serve":
+        create_parser(cli_args.make_arg_parser),
+        "chat":
+        create_parser(ChatCommand.add_cli_args),
+        "complete":
+        create_parser(CompleteCommand.add_cli_args),
+        "bench_latency":
+        create_parser(latency.add_cli_args),
+        "bench_throughput":
+        create_parser(throughput.add_cli_args),
+        "bench_serve":
+        create_parser(serve.add_cli_args),
+        "run-batch":
+        create_parser(run_batch.make_arg_parser),
     }
 
     # Generate documentation for each parser

docs/mkdocs/hooks/generate_examples.py

@@ -11,7 +11,7 @@ import regex as re
 logger = logging.getLogger("mkdocs")
 
 ROOT_DIR = Path(__file__).parent.parent.parent.parent
-ROOT_DIR_RELATIVE = "../../../../.."
+ROOT_DIR_RELATIVE = '../../../../..'
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/examples"
 
@@ -36,7 +36,7 @@ def fix_case(text: str) -> str:
         r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
     }
     for pattern, repl in subs.items():
-        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
+        text = re.sub(rf'\b{pattern}\b', repl, text, flags=re.IGNORECASE)
     return text
 
 
@@ -58,8 +58,7 @@ class Example:
         determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
         determine_title() -> str: Determines the title of the document.
         generate() -> str: Generates the documentation content.
-    """  # noqa: E501
-
+    """ # noqa: E501
     path: Path
     category: str = None
     main_file: Path = field(init=False)
@@ -85,8 +84,9 @@ class Example:
             Markdown file found in the directory.
         Raises:
             IndexError: If no Markdown files are found in the directory.
-        """  # noqa: E501
-        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
+        """ # noqa: E501
+        return self.path if self.path.is_file() else list(
+            self.path.glob("*.md")).pop()
 
     def determine_other_files(self) -> list[Path]:
         """
@@ -98,7 +98,7 @@ class Example:
 
         Returns:
             list[Path]: A list of Path objects representing the other files in the directory.
-        """  # noqa: E501
+        """ # noqa: E501
         if self.path.is_file():
             return []
         is_other_file = lambda file: file.is_file() and file != self.main_file
@@ -109,25 +109,25 @@ class Example:
             # Specify encoding for building on Windows
             with open(self.main_file, encoding="utf-8") as f:
                 first_line = f.readline().strip()
-            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            match = re.match(r'^#\s+(?P<title>.+)$', first_line)
             if match:
-                return match.group("title")
+                return match.group('title')
         return fix_case(self.path.stem.replace("_", " ").title())
 
     def fix_relative_links(self, content: str) -> str:
         """
         Fix relative links in markdown content by converting them to gh-file
         format.
-
+        
         Args:
             content (str): The markdown content to process
-
+            
         Returns:
             str: Content with relative links converted to gh-file format
         """
         # Regex to match markdown links [text](relative_path)
         # This matches links that don't start with http, https, ftp, or #
-        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
+        link_pattern = r'\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)'
 
         def replace_link(match):
             link_text = match.group(1)
@@ -137,7 +137,7 @@ class Example:
             gh_file = (self.main_file.parent / relative_path).resolve()
             gh_file = gh_file.relative_to(ROOT_DIR)
 
-            return f"[{link_text}](gh-file:{gh_file})"
+            return f'[{link_text}](gh-file:{gh_file})'
 
         return re.sub(link_pattern, replace_link, content)
 
@@ -150,11 +150,9 @@ class Example:
         code_fence = "``````"
 
         if self.is_code:
-            content += (
-                f"{code_fence}{self.main_file.suffix[1:]}\n"
-                f'--8<-- "{self.main_file}"\n'
-                f"{code_fence}\n"
-            )
+            content += (f"{code_fence}{self.main_file.suffix[1:]}\n"
+                        f'--8<-- "{self.main_file}"\n'
+                        f"{code_fence}\n")
         else:
             with open(self.main_file) as f:
                 # Skip the title from md snippets as it's been included above

docs/mkdocs/hooks/remove_announcement.py

@@ -7,7 +7,7 @@ from typing import Literal
 
 def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
     # see https://docs.readthedocs.io/en/stable/reference/environment-variables.html # noqa
-    if os.getenv("READTHEDOCS_VERSION_TYPE") == "tag":
+    if os.getenv('READTHEDOCS_VERSION_TYPE') == "tag":
         # remove the warning banner if the version is a tagged release
         mkdocs_dir = Path(__file__).parent.parent
         announcement_path = mkdocs_dir / "overrides/main.html"

docs/mkdocs/hooks/url_schemes.py

@@ -25,9 +25,8 @@ from mkdocs.structure.files import Files
 from mkdocs.structure.pages import Page
 
 
-def on_page_markdown(
-    markdown: str, *, page: Page, config: MkDocsConfig, files: Files
-) -> str:
+def on_page_markdown(markdown: str, *, page: Page, config: MkDocsConfig,
+                     files: Files) -> str:
     """
     Custom MkDocs plugin hook to rewrite special GitHub reference links
     in Markdown.
@@ -36,7 +35,7 @@ def on_page_markdown(
     GitHub shorthand links, such as:
       - `[Link text](gh-issue:123)`
       - `<gh-pr:456>`
-
+    
     And rewrites them into fully-qualified GitHub URLs with GitHub icons:
       - `[:octicons-mark-github-16: Link text](https://github.com/vllm-project/vllm/issues/123)`
       - `[:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)`
@@ -89,21 +88,21 @@ def on_page_markdown(
         """
         Replaces a matched inline-style GitHub shorthand link
         with a full Markdown link.
-
+        
         Example:
             [My issue](gh-issue:123) → [:octicons-mark-github-16: My issue](https://github.com/vllm-project/vllm/issues/123)
         """
-        url = f"{urls[match.group('type')]}/{match.group('path')}"
+        url = f'{urls[match.group("type")]}/{match.group("path")}'
         if fragment := match.group("fragment"):
             url += f"#{fragment}"
 
-        return f"[{gh_icon} {match.group('title')}]({url})"
+        return f'[{gh_icon} {match.group("title")}]({url})'
 
     def replace_auto_link(match: re.Match) -> str:
         """
         Replaces a matched autolink-style GitHub shorthand
         with a full Markdown link.
-
+        
         Example:
             <gh-pr:456> → [:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)
         """

docs/models/supported_models.md

@@ -32,9 +32,8 @@ If the Transformers model implementation follows all the steps in [writing a cus
 - All the features listed in the [compatibility matrix](../features/README.md#feature-x-feature)
 - Any combination of the following vLLM parallelisation schemes:
     - Data parallel
-    - Tensor parallel
-    - Expert parallel
     - Pipeline parallel
+    - Tensor parallel
 
 Checking if the modeling backend is Transformers is as simple as:
 
@@ -602,9 +601,8 @@ On the other hand, modalities separated by `/` are mutually exclusive.
 See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inputs to the model.
 
 !!! important
-    You can control the maximum number of multimodal inputs per prompt by setting
-    `limit_mm_per_prompt` (offline inference) or `--limit-mm-per-prompt` (online
-    serving). For example, to enable passing up to 4 images per text prompt:
+    **To enable multiple multi-modal items per text prompt in vLLM V0**, you have to set `limit_mm_per_prompt` (offline inference)
+    or `--limit-mm-per-prompt` (online serving). For example, to enable passing up to 4 images per text prompt:
 
     Offline inference:
 
@@ -623,6 +621,8 @@ See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inp
     vllm serve Qwen/Qwen2-VL-7B-Instruct --limit-mm-per-prompt '{"image":4}'
     ```
 
+    **This is no longer required if you are using vLLM V1.**
+
 !!! tip
     For hybrid-only models such as Llama-4, Step3 and Mistral-3, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (e.g, `--limit-mm-per-prompt '{"image":0}`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
 
@@ -730,7 +730,16 @@ Some models are supported only via the [Transformers backend](#transformers). Th
 <sup>+</sup> Multiple items can be inputted per text prompt for this modality.
 
 !!! warning
-    `Gemma3ForConditionalGeneration` uses a simplified attention pattern for text + image inputs:
+    Both V0 and V1 support `Gemma3ForConditionalGeneration` for text-only inputs.
+    However, there are differences in how they handle text + image inputs:
+
+    V0 correctly implements the model's attention pattern:
+    - Uses bidirectional attention between the image tokens corresponding to the same image
+    - Uses causal attention for other tokens
+    - Implemented via (naive) PyTorch SDPA with masking tensors
+    - Note: May use significant memory for long prompts with image
+
+    V1 currently uses a simplified attention pattern:
     - Uses causal attention for all tokens, including image tokens
     - Generates reasonable outputs but does not match the original model's attention for text + image inputs, especially when `{"do_pan_and_scan": true}`
     - Will be updated in the future to support the correct behavior
@@ -788,11 +797,11 @@ Some models are supported only via the [Transformers backend](#transformers). Th
     For more details, please see: <gh-pr:4087#issuecomment-2250397630>
 
 !!! warning
-    Our PaliGemma implementations currently share the same attention limitation as Gemma 3 (see above).
+    Our PaliGemma implementations have the same problem as Gemma 3 (see above) for both V0 and V1.
 
 !!! note
     For Qwen2.5-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`)
-    is currently unsupported because overlapping modalities are not yet supported.
+    is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1.
 
 #### Transcription
 
@@ -819,7 +828,6 @@ The following table lists those that are tested in vLLM.
 
 | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | [V1](gh-issue:8779) |
 |--------------|--------|--------|-------------------|----------------------|---------------------------|---------------------|
-| `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | | ✅︎ |
 | `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ | ✅︎ |
 | `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ | ✅︎ |
 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* | \* |

docs/usage/reproducibility.md

@@ -1,9 +1,10 @@
 # Reproducibility
 
-vLLM does not guarantee the reproducibility of the results by default, for the sake of performance. You need to do the following to achieve reproducible results:
+vLLM does not guarantee the reproducibility of the results by default, for the sake of performance. You need to do the following to achieve
+reproducible results:
 
-- Turn off multiprocessing to make the scheduling deterministic by setting `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
-- Optionally configure the global seed if you need to control random sampling (see below).
+- For V1: Turn off multiprocessing to make the scheduling deterministic by setting `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
+- For V0: Set the global seed (see below).
 
 Example: <gh-file:examples/offline_inference/reproducibility.py>
 
@@ -29,7 +30,9 @@ However, in some cases, setting the seed will also [change the random state in u
 
 ### Default Behavior
 
-The `seed` parameter defaults to `0`, which sets the random state for each worker so the results remain consistent for each vLLM run even if `temperature > 0`.
+In V0, the `seed` parameter defaults to `None`. When the `seed` parameter is `None`, the random states for `random`, `np.random`, and `torch.manual_seed` are not set. This means that each run of vLLM will produce different results if `temperature > 0`, as expected.
+
+In V1, the `seed` parameter defaults to `0` which sets the random state for each worker, so the results will remain consistent for each vLLM run even if `temperature > 0`.
 
 !!! note
 
@@ -40,6 +43,10 @@ The `seed` parameter defaults to `0`, which sets the random state for each worke
 
 ### Locality of random state
 
-The random state in user code (i.e. the code that constructs [LLM][vllm.LLM] class) is updated by vLLM when the workers run in the same process as user code, i.e.: `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
+The random state in user code (i.e. the code that constructs [LLM][vllm.LLM] class) is updated by vLLM under the following conditions:
 
-By default, this condition is not active so you can use vLLM without having to worry about accidentally making deterministic subsequent operations that rely on random state.
+- For V0: The seed is specified.
+- For V1: The workers are run in the same process as user code, i.e.: `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
+
+By default, these conditions are not active so you can use vLLM without having to worry about
+accidentally making deterministic subsequent operations that rely on random state.

docs/usage/v1_guide.md

@@ -1,16 +1,22 @@
 # vLLM V1
 
+!!! announcement
+
+    We have started the process of deprecating V0. Please read [RFC #18571](gh-issue:18571) for more details.
+
 V1 is now enabled by default for all supported use cases, and we will gradually enable it for every use case we plan to support. Please share any feedback on [GitHub](https://github.com/vllm-project/vllm) or in the [vLLM Slack](https://inviter.co/vllm-slack).
 
 To disable V1, please set the environment variable as: `VLLM_USE_V1=0`, and send us a GitHub issue sharing the reason!
 
 ## Why vLLM V1?
 
-vLLM V1 re-architects the engine to reduce accumulated complexity while preserving
-the stable, battle-tested components users rely on (such as models, GPU kernels,
-and supporting utilities). The scheduler, KV cache manager, worker, sampler, and
-API server now operate within a cohesive framework that is easier to extend and
-maintain as new capabilities are added.
+vLLM V0 successfully supported a wide range of models and hardware, but as new features were developed independently, the system grew increasingly complex. This complexity made it harder to integrate new capabilities and introduced technical debt, revealing the need for a more streamlined and unified design.
+
+Building on V0’s success, vLLM V1 retains the stable and proven components from V0
+(such as the models, GPU kernels, and utilities). At the same time, it significantly
+re-architects the core systems, covering the scheduler, KV cache manager, worker,
+sampler, and API server, to provide a cohesive, maintainable framework that better
+accommodates continued growth and innovation.
 
 Specifically, V1 aims to:
 
@@ -82,6 +88,8 @@ based on assigned priority, with FCFS as a tie-breaker), configurable via the
 | **Mamba Models**            | <nobr>🟢 (Mamba-2), 🟢 (Mamba-1)</nobr>                                            |
 | **Multimodal Models**       | <nobr>🟢 Functional</nobr>                                                         |
 
+vLLM V1 currently excludes model architectures with the `SupportsV0Only` protocol.
+
 !!! tip
 
     This corresponds to the V1 column in our [list of supported models](../models/supported_models.md).
@@ -141,8 +149,8 @@ encoder and decoder (e.g., `BartForConditionalGeneration`,
 
 #### Semantic Changes to Logprobs
 
-vLLM V1 supports logprobs and prompt logprobs. However, there are some important semantics
-to consider:
+vLLM V1 supports logprobs and prompt logprobs. However, there are some important semantic
+differences compared to V0:
 
 ##### Logprobs Calculation
 
@@ -167,7 +175,7 @@ As part of the major architectural rework in vLLM V1, several legacy features ha
 ##### Sampling features
 
 - **best_of**: This feature has been deprecated due to limited usage. See details at [RFC #13361](gh-issue:13361).
-- **Per-Request Logits Processors**: Previously, users could pass custom
+- **Per-Request Logits Processors**: In V0, users could pass custom
   processing functions to adjust logits on a per-request basis. In vLLM V1, this
   feature has been deprecated. Instead, the design is moving toward supporting **global logits
   processors**, a feature the team is actively working on for future releases. See details at [RFC #13360](gh-pr:13360).

examples/offline_inference/vision_language_multi_image.py

@@ -371,115 +371,6 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData:
     )
 
 
-def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=8192,
-        max_num_seqs=5,
-        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, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
-def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-1_5-8B"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=32768,
-        max_num_seqs=5,
-        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, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
-def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "moonshotai/Kimi-VL-A3B-Instruct"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=4096,
-        max_num_seqs=4,
-        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, trust_remote_code=True)
-
-    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_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
 
@@ -614,6 +505,115 @@ def load_llava_onevision(question: str, image_urls: list[str]) -> ModelRequestDa
     )
 
 
+def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:
+    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
+
+    engine_args = EngineArgs(
+        model=model_name,
+        trust_remote_code=True,
+        max_model_len=8192,
+        max_num_seqs=5,
+        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, trust_remote_code=True)
+
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    image_data = [fetch_image(url) for url in image_urls]
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompt=prompt,
+        image_data=image_data,
+    )
+
+
+def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData:
+    model_name = "Kwai-Keye/Keye-VL-1_5-8B"
+
+    engine_args = EngineArgs(
+        model=model_name,
+        trust_remote_code=True,
+        max_model_len=8192,
+        max_num_seqs=5,
+        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, trust_remote_code=True)
+
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    image_data = [fetch_image(url) for url in image_urls]
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompt=prompt,
+        image_data=image_data,
+    )
+
+
+def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
+    model_name = "moonshotai/Kimi-VL-A3B-Instruct"
+
+    engine_args = EngineArgs(
+        model=model_name,
+        trust_remote_code=True,
+        max_model_len=4096,
+        max_num_seqs=4,
+        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, trust_remote_code=True)
+
+    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_mistral3(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
 

examples/offline_inference/vision_language_pooling.py

@@ -58,30 +58,6 @@ class ModelRequestData(NamedTuple):
     documents: Optional[ScoreMultiModalParam] = None
 
 
-def run_clip(query: Query) -> ModelRequestData:
-    if query["modality"] == "text":
-        prompt = query["text"]
-        image = None
-    elif query["modality"] == "image":
-        prompt = ""  # For image input, make sure that the prompt text is empty
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: '{modality}'")
-
-    engine_args = EngineArgs(
-        model="openai/clip-vit-base-patch32",
-        runner="pooling",
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
 def run_e5_v(query: Query) -> ModelRequestData:
     llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n"  # noqa: E501
 
@@ -113,7 +89,7 @@ def run_e5_v(query: Query) -> ModelRequestData:
 def _get_vlm2vec_prompt_image(query: Query, image_token: str):
     if query["modality"] == "text":
         text = query["text"]
-        prompt = f"Find me an everyday image that matches the given caption: {text}"
+        prompt = f"Find me an everyday image that matches the given caption: {text}"  # noqa: E501
         image = None
     elif query["modality"] == "image":
         prompt = f"{image_token} Find a day-to-day image that looks similar to the provided image."  # noqa: E501
@@ -170,8 +146,7 @@ def run_vlm2vec_qwen2vl(query: Query) -> ModelRequestData:
 
     processor = AutoProcessor.from_pretrained(
         model_id,
-        # `min_pixels` and `max_pixels` are deprecated for
-        # transformers `preprocessor_config.json`
+        # `min_pixels` and `max_pixels` are deprecated
         size={"shortest_edge": 3136, "longest_edge": 12845056},
     )
     processor.chat_template = load_chat_template(
@@ -197,10 +172,8 @@ def run_vlm2vec_qwen2vl(query: Query) -> ModelRequestData:
         model=merged_path,
         runner="pooling",
         max_model_len=4096,
-        mm_processor_kwargs={
-            "min_pixels": 3136,
-            "max_pixels": 12845056,
-        },
+        trust_remote_code=True,
+        mm_processor_kwargs={"num_crops": 4},
         limit_mm_per_prompt={"image": 1},
     )
 
@@ -326,7 +299,6 @@ def run_score(model: str, modality: QueryModality, seed: Optional[int]):
 
 
 model_example_map = {
-    "clip": run_clip,
     "e5_v": run_e5_v,
     "vlm2vec_phi3v": run_vlm2vec_phi3v,
     "vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,

examples/online_serving/disaggregated_serving/disagg_proxy_demo.py

@@ -203,9 +203,9 @@ class Proxy:
                 async with session.post(
                     url=url, json=data, headers=headers
                 ) as response:
-                    if 200 <= response.status < 300 or 400 <= response.status < 500:
+                    if 200 <= response.status < 300 or 400 <= response.status < 500:  # noqa: E501
                         if use_chunked:
-                            async for chunk_bytes in response.content.iter_chunked(
+                            async for chunk_bytes in response.content.iter_chunked(  # noqa: E501
                                 1024
                             ):
                                 yield chunk_bytes

examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py

@@ -1,9 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""Example Python client for multimodal embedding API using vLLM API server.
-
-Refer to each `run_*` function for the command to run the server for that model.
+"""Example Python client for multimodal embedding API using vLLM API server
+NOTE:
+    start a supported multimodal embeddings model server with `vllm serve`, e.g.
+    vllm serve TIGER-Lab/VLM2Vec-Full \
+        --runner pooling \
+        --trust-remote-code \
+        --max-model-len 4096 \
+        --chat-template examples/template_vlm2vec_phi3v.jinja
 """
 
 import argparse
@@ -42,58 +47,7 @@ def create_chat_embeddings(
     )
 
 
-def run_clip(client: OpenAI, model: str):
-    """
-    Start the server using:
-
-    vllm serve openai/clip-vit-base-patch32 \
-        --runner pooling
-    """
-
-    response = create_chat_embeddings(
-        client,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": image_url}},
-                ],
-            }
-        ],
-        model=model,
-        encoding_format="float",
-    )
-
-    print("Image embedding output:", response.data[0].embedding)
-
-    response = create_chat_embeddings(
-        client,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": "a photo of a cat"},
-                ],
-            }
-        ],
-        model=model,
-        encoding_format="float",
-    )
-
-    print("Text embedding output:", response.data[0].embedding)
-
-
 def run_vlm2vec(client: OpenAI, model: str):
-    """
-    Start the server using:
-
-    vllm serve TIGER-Lab/VLM2Vec-Full \
-        --runner pooling \
-        --trust-remote-code \
-        --max-model-len 4096 \
-        --chat-template examples/template_vlm2vec_phi3v.jinja
-    """
-
     response = create_chat_embeddings(
         client,
         messages=[
@@ -149,15 +103,6 @@ def run_vlm2vec(client: OpenAI, model: str):
 
 
 def run_dse_qwen2_vl(client: OpenAI, model: str):
-    """
-    Start the server using:
-
-    vllm serve MrLight/dse-qwen2-2b-mrl-v1 \
-        --runner pooling \
-        --trust-remote-code \
-        --max-model-len 8192 \
-        --chat-template examples/template_dse_qwen2_vl.jinja
-    """
     response = create_chat_embeddings(
         client,
         messages=[
@@ -211,7 +156,6 @@ def run_dse_qwen2_vl(client: OpenAI, model: str):
 
 
 model_example_map = {
-    "clip": run_clip,
     "vlm2vec": run_vlm2vec,
     "dse_qwen2_vl": run_dse_qwen2_vl,
 }

examples/others/tensorize_vllm_model.py

@@ -21,6 +21,8 @@ from vllm.utils import FlexibleArgumentParser
 logger = logging.getLogger()
 
 
+# yapf conflicts with isort for this docstring
+# yapf: disable
 """
 tensorize_vllm_model.py is a script that can be used to serialize and 
 deserialize vLLM models. These models can be loaded using tensorizer 
@@ -130,8 +132,7 @@ def get_parser():
         "can be loaded using tensorizer directly to the GPU "
         "extremely quickly. Tensor encryption and decryption is "
         "also supported, although libsodium must be installed to "
-        "use it."
-    )
+        "use it.")
     parser = EngineArgs.add_cli_args(parser)
 
     parser.add_argument(
@@ -143,14 +144,13 @@ def get_parser():
         "along with the model by instantiating a TensorizerConfig object, "
         "creating a dict from it with TensorizerConfig.to_serializable(), "
         "and passing it to LoRARequest's initializer with the kwarg "
-        "tensorizer_config_dict.",
+        "tensorizer_config_dict."
     )
 
-    subparsers = parser.add_subparsers(dest="command", required=True)
+    subparsers = parser.add_subparsers(dest='command', required=True)
 
     serialize_parser = subparsers.add_parser(
-        "serialize", help="Serialize a model to `--serialized-directory`"
-    )
+        'serialize', help="Serialize a model to `--serialized-directory`")
 
     serialize_parser.add_argument(
         "--suffix",
@@ -163,9 +163,7 @@ def get_parser():
             "`--suffix` is `v1`, the serialized model tensors will be "
             "saved to "
             "`s3://my-bucket/vllm/EleutherAI/gpt-j-6B/v1/model.tensors`. "
-            "If none is provided, a random UUID will be used."
-        ),
-    )
+            "If none is provided, a random UUID will be used."))
     serialize_parser.add_argument(
         "--serialized-directory",
         type=str,
@@ -177,127 +175,108 @@ def get_parser():
         "and the model HuggingFace ID is `EleutherAI/gpt-j-6B`, tensors will "
         "be saved to `dir/vllm/EleutherAI/gpt-j-6B/suffix/model.tensors`, "
         "where `suffix` is given by `--suffix` or a random UUID if not "
-        "provided.",
-    )
+        "provided.")
 
     serialize_parser.add_argument(
         "--serialization-kwargs",
         type=tensorizer_kwargs_arg,
         required=False,
-        help=(
-            "A JSON string containing additional keyword arguments to "
-            "pass to Tensorizer's TensorSerializer during "
-            "serialization."
-        ),
-    )
+        help=("A JSON string containing additional keyword arguments to "
+              "pass to Tensorizer's TensorSerializer during "
+              "serialization."))
 
     serialize_parser.add_argument(
         "--keyfile",
         type=str,
         required=False,
-        help=(
-            "Encrypt the model weights with a randomly-generated binary key,"
-            " and save the key at this path"
-        ),
-    )
+        help=("Encrypt the model weights with a randomly-generated binary key,"
+              " and save the key at this path"))
 
     deserialize_parser = subparsers.add_parser(
-        "deserialize",
-        help=(
-            "Deserialize a model from `--path-to-tensors`"
-            " to verify it can be loaded and used."
-        ),
-    )
+        'deserialize',
+        help=("Deserialize a model from `--path-to-tensors`"
+              " to verify it can be loaded and used."))
 
     deserialize_parser.add_argument(
         "--path-to-tensors",
         type=str,
         required=False,
-        help="The local path or S3 URI to the model tensors to deserialize. ",
-    )
+        help="The local path or S3 URI to the model tensors to deserialize. ")
 
     deserialize_parser.add_argument(
         "--serialized-directory",
         type=str,
         required=False,
         help="Directory with model artifacts for loading. Assumes a "
-        "model.tensors file exists therein. Can supersede "
-        "--path-to-tensors.",
-    )
+             "model.tensors file exists therein. Can supersede "
+             "--path-to-tensors.")
 
     deserialize_parser.add_argument(
         "--keyfile",
         type=str,
         required=False,
-        help=(
-            "Path to a binary key to use to decrypt the model weights,"
-            " if the model was serialized with encryption"
-        ),
-    )
+        help=("Path to a binary key to use to decrypt the model weights,"
+              " if the model was serialized with encryption"))
 
     deserialize_parser.add_argument(
         "--deserialization-kwargs",
         type=tensorizer_kwargs_arg,
         required=False,
-        help=(
-            "A JSON string containing additional keyword arguments to "
-            "pass to Tensorizer's `TensorDeserializer` during "
-            "deserialization."
-        ),
-    )
+        help=("A JSON string containing additional keyword arguments to "
+              "pass to Tensorizer's `TensorDeserializer` during "
+              "deserialization."))
 
     TensorizerArgs.add_cli_args(deserialize_parser)
 
     return parser
 
-
-def merge_extra_config_with_tensorizer_config(extra_cfg: dict, cfg: TensorizerConfig):
+def merge_extra_config_with_tensorizer_config(extra_cfg: dict,
+                                              cfg: TensorizerConfig):
     for k, v in extra_cfg.items():
         if hasattr(cfg, k):
             setattr(cfg, k, v)
             logger.info(
                 "Updating TensorizerConfig with %s from "
-                "--model-loader-extra-config provided",
-                k,
+                "--model-loader-extra-config provided", k
             )
 
-
 def deserialize(args, tensorizer_config):
     if args.lora_path:
         tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
-        llm = LLM(
-            model=args.model,
-            load_format="tensorizer",
-            tensor_parallel_size=args.tensor_parallel_size,
-            model_loader_extra_config=tensorizer_config,
-            enable_lora=True,
+        llm = LLM(model=args.model,
+                  load_format="tensorizer",
+                  tensor_parallel_size=args.tensor_parallel_size,
+                  model_loader_extra_config=tensorizer_config,
+                  enable_lora=True,
         )
         sampling_params = SamplingParams(
-            temperature=0, max_tokens=256, stop=["[/assistant]"]
+            temperature=0,
+            max_tokens=256,
+            stop=["[/assistant]"]
         )
 
         # Truncating this as the extra text isn't necessary
-        prompts = ["[user] Write a SQL query to answer the question based on ..."]
+        prompts = [
+            "[user] Write a SQL query to answer the question based on ..."
+        ]
 
         # Test LoRA load
         print(
             llm.generate(
-                prompts,
-                sampling_params,
-                lora_request=LoRARequest(
-                    "sql-lora",
-                    1,
-                    args.lora_path,
-                    tensorizer_config_dict=tensorizer_config.to_serializable(),
-                ),
+            prompts,
+            sampling_params,
+            lora_request=LoRARequest("sql-lora",
+                                     1,
+                                     args.lora_path,
+                                     tensorizer_config_dict = tensorizer_config
+                                     .to_serializable())
             )
         )
     else:
-        llm = LLM(
-            model=args.model,
-            load_format="tensorizer",
-            tensor_parallel_size=args.tensor_parallel_size,
-            model_loader_extra_config=tensorizer_config,
+        llm = LLM(model=args.model,
+                  load_format="tensorizer",
+                  tensor_parallel_size=args.tensor_parallel_size,
+                  model_loader_extra_config=tensorizer_config
         )
     return llm
 
@@ -306,20 +285,17 @@ def main():
     parser = get_parser()
     args = parser.parse_args()
 
-    s3_access_key_id = getattr(args, "s3_access_key_id", None) or os.environ.get(
-        "S3_ACCESS_KEY_ID", None
-    )
-    s3_secret_access_key = getattr(
-        args, "s3_secret_access_key", None
-    ) or os.environ.get("S3_SECRET_ACCESS_KEY", None)
-    s3_endpoint = getattr(args, "s3_endpoint", None) or os.environ.get(
-        "S3_ENDPOINT_URL", None
-    )
+    s3_access_key_id = (getattr(args, 's3_access_key_id', None)
+                        or os.environ.get("S3_ACCESS_KEY_ID", None))
+    s3_secret_access_key = (getattr(args, 's3_secret_access_key', None)
+                            or os.environ.get("S3_SECRET_ACCESS_KEY", None))
+    s3_endpoint = (getattr(args, 's3_endpoint', None)
+                or os.environ.get("S3_ENDPOINT_URL", None))
 
     credentials = {
         "s3_access_key_id": s3_access_key_id,
         "s3_secret_access_key": s3_secret_access_key,
-        "s3_endpoint": s3_endpoint,
+        "s3_endpoint": s3_endpoint
     }
 
     model_ref = args.model
@@ -333,25 +309,25 @@ def main():
     if args.model_loader_extra_config:
         extra_config = json.loads(args.model_loader_extra_config)
 
-    tensorizer_dir = args.serialized_directory or extra_config.get("tensorizer_dir")
-    tensorizer_uri = getattr(args, "path_to_tensors", None) or extra_config.get(
-        "tensorizer_uri"
-    )
+
+    tensorizer_dir = (args.serialized_directory or
+                      extra_config.get("tensorizer_dir"))
+    tensorizer_uri = (getattr(args, "path_to_tensors", None)
+                      or extra_config.get("tensorizer_uri"))
 
     if tensorizer_dir and tensorizer_uri:
-        parser.error(
-            "--serialized-directory and --path-to-tensors cannot both be provided"
-        )
+        parser.error("--serialized-directory and --path-to-tensors "
+                     "cannot both be provided")
 
     if not tensorizer_dir and not tensorizer_uri:
-        parser.error(
-            "Either --serialized-directory or --path-to-tensors must be provided"
-        )
+        parser.error("Either --serialized-directory or --path-to-tensors "
+                     "must be provided")
+
 
     if args.command == "serialize":
         engine_args = EngineArgs.from_cli_args(args)
 
-        input_dir = tensorizer_dir.rstrip("/")
+        input_dir = tensorizer_dir.rstrip('/')
         suffix = args.suffix if args.suffix else uuid.uuid4().hex
         base_path = f"{input_dir}/vllm/{model_ref}/{suffix}"
         if engine_args.tensor_parallel_size > 1:
@@ -363,14 +339,15 @@ def main():
             tensorizer_uri=model_path,
             encryption_keyfile=keyfile,
             serialization_kwargs=args.serialization_kwargs or {},
-            **credentials,
+            **credentials
         )
 
         if args.lora_path:
             tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
             tensorize_lora_adapter(args.lora_path, tensorizer_config)
 
-        merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config)
+        merge_extra_config_with_tensorizer_config(extra_config,
+                                                  tensorizer_config)
         tensorize_vllm_model(engine_args, tensorizer_config)
 
     elif args.command == "deserialize":
@@ -379,10 +356,11 @@ def main():
             tensorizer_dir=args.serialized_directory,
             encryption_keyfile=keyfile,
             deserialization_kwargs=args.deserialization_kwargs or {},
-            **credentials,
+            **credentials
         )
 
-        merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config)
+        merge_extra_config_with_tensorizer_config(extra_config,
+                                                  tensorizer_config)
         deserialize(args, tensorizer_config)
     else:
         raise ValueError("Either serialize or deserialize must be specified.")

examples/pyproject.toml

@@ -0,0 +1,54 @@
+# This local pyproject file is part of the migration from yapf to ruff format.
+# It uses the same core rules as the main pyproject.toml file, but with the
+# following differences:
+# - ruff line length is overridden to 88
+# - deprecated typing ignores (UP006, UP035) have been removed
+
+[tool.ruff]
+line-length = 88
+exclude = [
+    # External file, leaving license intact
+    "examples/other/fp8/quantizer/quantize.py",
+    "vllm/vllm_flash_attn/flash_attn_interface.pyi"
+]
+
+[tool.ruff.lint.per-file-ignores]
+"vllm/third_party/**" = ["ALL"]
+"vllm/version.py" = ["F401"]
+"vllm/_version.py" = ["ALL"]
+
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    "I",
+    # flake8-logging-format
+    "G",
+]
+ignore = [
+    # star imports
+    "F405", "F403",
+    # lambda expression assignment
+    "E731",
+    # Loop control variable not used within loop body
+    "B007",
+    # f-string format
+    "UP032",
+    # Can remove once 3.10+ is the minimum Python version
+    "UP007",
+]
+
+[tool.ruff.lint.isort]
+known-first-party = ["vllm"]
+
+[tool.ruff.format]
+docstring-code-format = true

pyproject.toml

@@ -52,10 +52,27 @@ lora_filesystem_resolver = "vllm.plugins.lora_resolvers.filesystem_resolver:regi
 where = ["."]
 include = ["vllm*"]
 
+[tool.yapfignore]
+ignore_patterns = [
+    ".buildkite/**",
+    "benchmarks/**",
+    "build/**",
+    "examples/**",
+]
+
+[tool.ruff]
+# Allow lines to be as long as 80.
+line-length = 80
+
 [tool.ruff.lint.per-file-ignores]
 "vllm/third_party/**" = ["ALL"]
 "vllm/version.py" = ["F401"]
 "vllm/_version.py" = ["ALL"]
+# Python 3.8 typing - skip V0 code
+"vllm/attention/**/*.py" = ["UP006", "UP035"]
+"vllm/engine/**/*.py" = ["UP006", "UP035"]
+"vllm/executor/**/*.py" = ["UP006", "UP035"]
+"vllm/worker/**/*.py" = ["UP006", "UP035"]
 
 [tool.ruff.lint]
 select = [
@@ -70,7 +87,7 @@ select = [
     # flake8-simplify
     "SIM",
     # isort
-    "I",
+    # "I",
     # flake8-logging-format
     "G",
 ]
@@ -87,15 +104,21 @@ ignore = [
     "UP007",
 ]
 
-[tool.ruff.format]
-docstring-code-format = true
-
 [tool.mypy]
 plugins = ['pydantic.mypy']
 ignore_missing_imports = true
 check_untyped_defs = true
 follow_imports = "silent"
 
+[tool.isort]
+skip_glob = [
+    ".buildkite/*",
+    "benchmarks/*",
+    "examples/*",
+]
+use_parentheses = true
+skip_gitignore = true
+
 [tool.pytest.ini_options]
 markers = [
     "slow_test",

requirements/common.txt

@@ -49,4 +49,3 @@ pybase64 # fast base64 implementation
 cbor2 # Required for cross-language serialization of hashable objects
 setproctitle # Used to set process names for better debugging and monitoring
 openai-harmony >= 0.0.3  # Required for gpt-oss
-gpt-oss >= 0.0.7

setup.py

@@ -34,36 +34,32 @@ logger = logging.getLogger(__name__)
 
 # cannot import envs directly because it depends on vllm,
 #  which is not installed yet
-envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
+envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'vllm', 'envs.py'))
 
 VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
 
 if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
-    logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
+    logger.warning(
+        "VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
     VLLM_TARGET_DEVICE = "cpu"
-elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")):
+elif not (sys.platform.startswith("linux")
+          or sys.platform.startswith("darwin")):
     logger.warning(
         "vLLM only supports Linux platform (including WSL) and MacOS."
         "Building on %s, "
-        "so vLLM may not be able to run correctly",
-        sys.platform,
-    )
+        "so vLLM may not be able to run correctly", sys.platform)
     VLLM_TARGET_DEVICE = "empty"
-elif (
-    sys.platform.startswith("linux")
-    and torch.version.cuda is None
-    and os.getenv("VLLM_TARGET_DEVICE") is None
-    and torch.version.hip is None
-):
+elif (sys.platform.startswith("linux") and torch.version.cuda is None
+      and os.getenv("VLLM_TARGET_DEVICE") is None
+      and torch.version.hip is None):
     # if cuda or hip is not available and VLLM_TARGET_DEVICE is not set,
     # fallback to cpu
     VLLM_TARGET_DEVICE = "cpu"
 
 
 def is_sccache_available() -> bool:
-    return which("sccache") is not None and not bool(
-        int(os.getenv("VLLM_DISABLE_SCCACHE", "0"))
-    )
+    return which("sccache") is not None and \
+        not bool(int(os.getenv("VLLM_DISABLE_SCCACHE", "0")))
 
 
 def is_ccache_available() -> bool:
@@ -87,7 +83,8 @@ def is_url_available(url: str) -> bool:
 
 
 class CMakeExtension(Extension):
-    def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
+
+    def __init__(self, name: str, cmake_lists_dir: str = '.', **kwa) -> None:
         super().__init__(name, sources=[], py_limited_api=True, **kwa)
         self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)
 
@@ -124,8 +121,8 @@ class cmake_build_ext(build_ext):
             if nvcc_threads is not None:
                 nvcc_threads = int(nvcc_threads)
                 logger.info(
-                    "Using NVCC_THREADS=%d as the number of nvcc threads.", nvcc_threads
-                )
+                    "Using NVCC_THREADS=%d as the number of nvcc threads.",
+                    nvcc_threads)
             else:
                 nvcc_threads = 1
             num_jobs = max(1, num_jobs // nvcc_threads)
@@ -149,36 +146,36 @@ class cmake_build_ext(build_ext):
         cfg = envs.CMAKE_BUILD_TYPE or default_cfg
 
         cmake_args = [
-            "-DCMAKE_BUILD_TYPE={}".format(cfg),
-            "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
+            '-DCMAKE_BUILD_TYPE={}'.format(cfg),
+            '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
         ]
 
         verbose = envs.VERBOSE
         if verbose:
-            cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"]
+            cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']
 
         if is_sccache_available():
             cmake_args += [
-                "-DCMAKE_C_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache",
-                "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache",
+                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
+                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
+                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
+                '-DCMAKE_HIP_COMPILER_LAUNCHER=sccache',
             ]
         elif is_ccache_available():
             cmake_args += [
-                "-DCMAKE_C_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache",
-                "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache",
+                '-DCMAKE_C_COMPILER_LAUNCHER=ccache',
+                '-DCMAKE_CXX_COMPILER_LAUNCHER=ccache',
+                '-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache',
+                '-DCMAKE_HIP_COMPILER_LAUNCHER=ccache',
             ]
 
         # Pass the python executable to cmake so it can find an exact
         # match.
-        cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)]
+        cmake_args += ['-DVLLM_PYTHON_EXECUTABLE={}'.format(sys.executable)]
 
         # Pass the python path to cmake so it can reuse the build dependencies
         # on subsequent calls to python.
-        cmake_args += ["-DVLLM_PYTHON_PATH={}".format(":".join(sys.path))]
+        cmake_args += ['-DVLLM_PYTHON_PATH={}'.format(":".join(sys.path))]
 
         # Override the base directory for FetchContent downloads to $ROOT/.deps
         # This allows sharing dependencies between profiles,
@@ -186,7 +183,7 @@ class cmake_build_ext(build_ext):
         # To override this, set the FETCHCONTENT_BASE_DIR environment variable.
         fc_base_dir = os.path.join(ROOT_DIR, ".deps")
         fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir)
-        cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)]
+        cmake_args += ['-DFETCHCONTENT_BASE_DIR={}'.format(fc_base_dir)]
 
         #
         # Setup parallelism and build tool
@@ -194,36 +191,30 @@ class cmake_build_ext(build_ext):
         num_jobs, nvcc_threads = self.compute_num_jobs()
 
         if nvcc_threads:
-            cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)]
+            cmake_args += ['-DNVCC_THREADS={}'.format(nvcc_threads)]
 
         if is_ninja_available():
-            build_tool = ["-G", "Ninja"]
+            build_tool = ['-G', 'Ninja']
             cmake_args += [
-                "-DCMAKE_JOB_POOL_COMPILE:STRING=compile",
-                "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs),
+                '-DCMAKE_JOB_POOL_COMPILE:STRING=compile',
+                '-DCMAKE_JOB_POOLS:STRING=compile={}'.format(num_jobs),
             ]
         else:
             # Default build tool to whatever cmake picks.
             build_tool = []
         # Make sure we use the nvcc from CUDA_HOME
         if _is_cuda():
-            cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"]
-
-        other_cmake_args = os.environ.get("CMAKE_ARGS")
-        if other_cmake_args:
-            cmake_args += other_cmake_args.split()
-
+            cmake_args += [f'-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc']
         subprocess.check_call(
-            ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args],
-            cwd=self.build_temp,
-        )
+            ['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args],
+            cwd=self.build_temp)
 
     def build_extensions(self) -> None:
         # Ensure that CMake is present and working
         try:
-            subprocess.check_output(["cmake", "--version"])
+            subprocess.check_output(['cmake', '--version'])
         except OSError as e:
-            raise RuntimeError("Cannot find CMake executable") from e
+            raise RuntimeError('Cannot find CMake executable') from e
 
         # Create build directory if it does not exist.
         if not os.path.exists(self.build_temp):
@@ -262,18 +253,13 @@ class cmake_build_ext(build_ext):
             # CMake appends the extension prefix to the install path,
             # and outdir already contains that prefix, so we need to remove it.
             prefix = outdir
-            for _ in range(ext.name.count(".")):
+            for _ in range(ext.name.count('.')):
                 prefix = prefix.parent
 
             # prefix here should actually be the same for all components
             install_args = [
-                "cmake",
-                "--install",
-                ".",
-                "--prefix",
-                prefix,
-                "--component",
-                target_name(ext.name),
+                "cmake", "--install", ".", "--prefix", prefix, "--component",
+                target_name(ext.name)
             ]
             subprocess.check_call(install_args, cwd=self.build_temp)
 
@@ -284,15 +270,12 @@ class cmake_build_ext(build_ext):
         # copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current
         # directory so that they can be included in the editable build
         import glob
-
-        files = glob.glob(
-            os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "**", "*.py"),
-            recursive=True,
-        )
+        files = glob.glob(os.path.join(self.build_lib, "vllm",
+                                       "vllm_flash_attn", "**", "*.py"),
+                          recursive=True)
         for file in files:
-            dst_file = os.path.join(
-                "vllm/vllm_flash_attn", file.split("vllm/vllm_flash_attn/")[-1]
-            )
+            dst_file = os.path.join("vllm/vllm_flash_attn",
+                                    file.split("vllm/vllm_flash_attn/")[-1])
             print(f"Copying {file} to {dst_file}")
             os.makedirs(os.path.dirname(dst_file), exist_ok=True)
             self.copy_file(file, dst_file)
@@ -302,7 +285,8 @@ class precompiled_build_ext(build_ext):
     """Disables extension building when using precompiled binaries."""
 
     def run(self) -> None:
-        assert _is_cuda(), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
+        assert _is_cuda(
+        ), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
 
     def build_extensions(self) -> None:
         print("Skipping build_ext: using precompiled extensions.")
@@ -323,9 +307,9 @@ class precompiled_wheel_utils:
                 wheel_filename = wheel_url_or_path.split("/")[-1]
                 temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
                 wheel_path = os.path.join(temp_dir, wheel_filename)
-                print(f"Downloading wheel from {wheel_url_or_path} to {wheel_path}")
+                print(f"Downloading wheel from {wheel_url_or_path} "
+                      f"to {wheel_path}")
                 from urllib.request import urlretrieve
-
                 urlretrieve(wheel_url_or_path, filename=wheel_path)
             else:
                 wheel_path = wheel_url_or_path
@@ -346,29 +330,25 @@ class precompiled_wheel_utils:
                 ]
 
                 compiled_regex = re.compile(
-                    r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
-                )
+                    r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
                 file_members = list(
-                    filter(lambda x: x.filename in files_to_copy, wheel.filelist)
-                )
+                    filter(lambda x: x.filename in files_to_copy,
+                           wheel.filelist))
                 file_members += list(
-                    filter(lambda x: compiled_regex.match(x.filename), wheel.filelist)
-                )
+                    filter(lambda x: compiled_regex.match(x.filename),
+                           wheel.filelist))
 
                 for file in file_members:
                     print(f"[extract] {file.filename}")
                     target_path = os.path.join(".", file.filename)
                     os.makedirs(os.path.dirname(target_path), exist_ok=True)
-                    with (
-                        wheel.open(file.filename) as src,
-                        open(target_path, "wb") as dst,
-                    ):
+                    with wheel.open(file.filename) as src, open(
+                            target_path, "wb") as dst:
                         shutil.copyfileobj(src, dst)
 
                     pkg = os.path.dirname(file.filename).replace("/", ".")
                     package_data_patch.setdefault(pkg, []).append(
-                        os.path.basename(file.filename)
-                    )
+                        os.path.basename(file.filename))
 
             return package_data_patch
         finally:
@@ -384,13 +364,10 @@ class precompiled_wheel_utils:
 
         try:
             # Get the latest commit hash of the upstream main branch.
-            resp_json = subprocess.check_output(
-                [
-                    "curl",
-                    "-s",
-                    "https://api.github.com/repos/vllm-project/vllm/commits/main",
-                ]
-            ).decode("utf-8")
+            resp_json = subprocess.check_output([
+                "curl", "-s",
+                "https://api.github.com/repos/vllm-project/vllm/commits/main"
+            ]).decode("utf-8")
             upstream_main_commit = json.loads(resp_json)["sha"]
 
             # In Docker build context, .git may be immutable or missing.
@@ -400,32 +377,25 @@ class precompiled_wheel_utils:
             # Check if the upstream_main_commit exists in the local repo
             try:
                 subprocess.check_output(
-                    ["git", "cat-file", "-e", f"{upstream_main_commit}"]
-                )
+                    ["git", "cat-file", "-e", f"{upstream_main_commit}"])
             except subprocess.CalledProcessError:
                 # If not present, fetch it from the remote repository.
                 # Note that this does not update any local branches,
                 # but ensures that this commit ref and its history are
                 # available in our local repo.
-                subprocess.check_call(
-                    ["git", "fetch", "https://github.com/vllm-project/vllm", "main"]
-                )
+                subprocess.check_call([
+                    "git", "fetch", "https://github.com/vllm-project/vllm",
+                    "main"
+                ])
 
             # Then get the commit hash of the current branch that is the same as
             # the upstream main commit.
-            current_branch = (
-                subprocess.check_output(["git", "branch", "--show-current"])
-                .decode("utf-8")
-                .strip()
-            )
+            current_branch = subprocess.check_output(
+                ["git", "branch", "--show-current"]).decode("utf-8").strip()
 
-            base_commit = (
-                subprocess.check_output(
-                    ["git", "merge-base", f"{upstream_main_commit}", current_branch]
-                )
-                .decode("utf-8")
-                .strip()
-            )
+            base_commit = subprocess.check_output([
+                "git", "merge-base", f"{upstream_main_commit}", current_branch
+            ]).decode("utf-8").strip()
             return base_commit
         except ValueError as err:
             raise ValueError(err) from None
@@ -433,9 +403,7 @@ class precompiled_wheel_utils:
             logger.warning(
                 "Failed to get the base commit in the main branch. "
                 "Using the nightly wheel. The libraries in this "
-                "wheel may not be compatible with your dev branch: %s",
-                err,
-            )
+                "wheel may not be compatible with your dev branch: %s", err)
             return "nightly"
 
 
@@ -445,13 +413,12 @@ def _no_device() -> bool:
 
 def _is_cuda() -> bool:
     has_cuda = torch.version.cuda is not None
-    return VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu()
+    return (VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu())
 
 
 def _is_hip() -> bool:
-    return (
-        VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm"
-    ) and torch.version.hip is not None
+    return (VLLM_TARGET_DEVICE == "cuda"
+            or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
 
 
 def _is_tpu() -> bool:
@@ -490,12 +457,8 @@ def get_rocm_version():
         minor = ctypes.c_uint32()
         patch = ctypes.c_uint32()
 
-        if (
-            get_rocm_core_version(
-                ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
-            )
-            == 0
-        ):
+        if (get_rocm_core_version(ctypes.byref(major), ctypes.byref(minor),
+                                  ctypes.byref(patch)) == 0):
             return f"{major.value}.{minor.value}.{patch.value}"
         return None
     except Exception:
@@ -508,9 +471,8 @@ def get_nvcc_cuda_version() -> Version:
     Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
     """
     assert CUDA_HOME is not None, "CUDA_HOME is not set"
-    nvcc_output = subprocess.check_output(
-        [CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True
-    )
+    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
+                                          universal_newlines=True)
     output = nvcc_output.split()
     release_idx = output.index("release") + 1
     nvcc_cuda_version = parse(output[release_idx].split(",")[0])
@@ -522,20 +484,14 @@ def get_gaudi_sw_version():
     Returns the driver version.
     """
     # Enable console printing for `hl-smi` check
-    output = subprocess.run(
-        "hl-smi",
-        shell=True,
-        text=True,
-        capture_output=True,
-        env={"ENABLE_CONSOLE": "true"},
-    )
+    output = subprocess.run("hl-smi",
+                            shell=True,
+                            text=True,
+                            capture_output=True,
+                            env={"ENABLE_CONSOLE": "true"})
     if output.returncode == 0 and output.stdout:
-        return (
-            output.stdout.split("\n")[2]
-            .replace(" ", "")
-            .split(":")[1][:-1]
-            .split("-")[0]
-        )
+        return output.stdout.split("\n")[2].replace(
+            " ", "").split(":")[1][:-1].split("-")[0]
     return "0.0.0"  # when hl-smi is not available
 
 
@@ -585,11 +541,8 @@ def get_requirements() -> list[str]:
         for line in requirements:
             if line.startswith("-r "):
                 resolved_requirements += _read_requirements(line.split()[1])
-            elif (
-                not line.startswith("--")
-                and not line.startswith("#")
-                and line.strip() != ""
-            ):
+            elif not line.startswith("--") and not line.startswith(
+                    "#") and line.strip() != "":
                 resolved_requirements.append(line)
         return resolved_requirements
 
@@ -600,7 +553,7 @@ def get_requirements() -> list[str]:
         cuda_major, cuda_minor = torch.version.cuda.split(".")
         modified_requirements = []
         for req in requirements:
-            if "vllm-flash-attn" in req and cuda_major != "12":
+            if ("vllm-flash-attn" in req and cuda_major != "12"):
                 # vllm-flash-attn is built only for CUDA 12.x.
                 # Skip for other versions.
                 continue
@@ -615,7 +568,8 @@ def get_requirements() -> list[str]:
     elif _is_xpu():
         requirements = _read_requirements("xpu.txt")
     else:
-        raise ValueError("Unsupported platform, please use CUDA, ROCm, or CPU.")
+        raise ValueError(
+            "Unsupported platform, please use CUDA, ROCm, or CPU.")
     return requirements
 
 
@@ -631,13 +585,14 @@ if _is_cuda():
     ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
     if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
         # FA3 requires CUDA 12.3 or later
-        ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
+        ext_modules.append(
+            CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
         # Optional since this doesn't get built (produce an .so file) when
         # not targeting a hopper system
-        ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
         ext_modules.append(
-            CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
-        )
+            CMakeExtension(name="vllm._flashmla_C", optional=True))
+        ext_modules.append(
+            CMakeExtension(name="vllm._flashmla_extension_C", optional=True))
     ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
 
 if _build_custom_ops():
@@ -659,7 +614,6 @@ if envs.VLLM_USE_PRECOMPILED:
         wheel_url = wheel_location
     else:
         import platform
-
         arch = platform.machine()
         if arch == "x86_64":
             wheel_tag = "manylinux1_x86_64"
@@ -669,11 +623,8 @@ if envs.VLLM_USE_PRECOMPILED:
             raise ValueError(f"Unsupported architecture: {arch}")
         base_commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
         wheel_url = f"https://wheels.vllm.ai/{base_commit}/vllm-1.0.0.dev-cp38-abi3-{wheel_tag}.whl"
-        nightly_wheel_url = (
-            f"https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-{wheel_tag}.whl"
-        )
+        nightly_wheel_url = f"https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-{wheel_tag}.whl"
         from urllib.request import urlopen
-
         try:
             with urlopen(wheel_url) as resp:
                 if resp.status != 200:
@@ -682,7 +633,8 @@ if envs.VLLM_USE_PRECOMPILED:
             print(f"[warn] Falling back to nightly wheel: {e}")
             wheel_url = nightly_wheel_url
 
-    patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(wheel_url)
+    patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(
+        wheel_url)
     for pkg, files in patch.items():
         package_data.setdefault(pkg, []).extend(files)
 
@@ -693,9 +645,8 @@ if not ext_modules:
     cmdclass = {}
 else:
     cmdclass = {
-        "build_ext": precompiled_build_ext
-        if envs.VLLM_USE_PRECOMPILED
-        else cmake_build_ext
+        "build_ext":
+        precompiled_build_ext if envs.VLLM_USE_PRECOMPILED else cmake_build_ext
     }
 
 setup(
@@ -708,11 +659,8 @@ setup(
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.1.10"],
         "runai": ["runai-model-streamer[s3,gcs] >= 0.14.0"],
-        "audio": [
-            "librosa",
-            "soundfile",
-            "mistral_common[audio]",
-        ],  # Required for audio processing
+        "audio": ["librosa", "soundfile",
+                  "mistral_common[audio]"],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
         # FlashInfer should be updated together with the Dockerfile
         "flashinfer": ["flashinfer-python==0.3.1"],

tests/basic_correctness/test_basic_correctness.py

@@ -4,7 +4,6 @@
 
 Run `pytest tests/basic_correctness/test_basic_correctness.py`.
 """
-
 import os
 import weakref
 from unittest.mock import Mock
@@ -38,21 +37,16 @@ def test_vllm_gc_ed():
 
 
 def _fix_prompt_embed_outputs(
-    vllm_outputs: list[tuple[list[int], str]],
-    hf_model: HfRunner,
-    example_prompts: list[str],
-) -> list[tuple[list[int], str]]:
+        vllm_outputs: list[tuple[list[int], str]], hf_model: HfRunner,
+        example_prompts: list[str]) -> list[tuple[list[int], str]]:
     fixed_vllm_outputs = []
     for vllm_output, hf_input, prompt in zip(
-        vllm_outputs, hf_model.get_inputs(example_prompts), example_prompts
-    ):
+            vllm_outputs, hf_model.get_inputs(example_prompts),
+            example_prompts):
         hf_input_ids = hf_input["input_ids"].tolist()[0]
         fixed_vllm_outputs.append(
-            (
-                hf_input_ids + vllm_output[0][len(hf_input_ids) :],
-                prompt + vllm_output[1],
-            )
-        )
+            (hf_input_ids + vllm_output[0][len(hf_input_ids):],
+             prompt + vllm_output[1]))
     return fixed_vllm_outputs
 
 
@@ -75,7 +69,8 @@ def test_models(
     enable_prompt_embeds: bool,
 ) -> None:
     if backend == "XFORMERS" and model == "google/gemma-2-2b-it":
-        pytest.skip(f"{backend} does not support gemma2 with full context length.")
+        pytest.skip(
+            f"{backend} does not support gemma2 with full context length.")
 
     with monkeypatch.context() as m:
         m.setenv("VLLM_ATTENTION_BACKEND", backend)
@@ -83,35 +78,34 @@ def test_models(
         # 5042 tokens for gemma2
         # gemma2 has alternating sliding window size of 4096
         # we need a prompt with more than 4096 tokens to test the sliding window
-        prompt = (
-            "The following numbers of the sequence "
-            + ", ".join(str(i) for i in range(1024))
-            + " are:"
-        )
+        prompt = "The following numbers of the sequence " + ", ".join(
+            str(i) for i in range(1024)) + " are:"
         example_prompts = [prompt]
 
         with hf_runner(model) as hf_model:
             hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
             if enable_prompt_embeds:
                 with torch.no_grad():
-                    prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
+                    prompt_embeds = hf_model.get_prompt_embeddings(
+                        example_prompts)
 
         with VllmRunner(
-            model,
-            max_model_len=8192,
-            enforce_eager=enforce_eager,
-            enable_prompt_embeds=enable_prompt_embeds,
-            gpu_memory_utilization=0.7,
-            async_scheduling=async_scheduling,
-            distributed_executor_backend=model_executor,
+                model,
+                max_model_len=8192,
+                enforce_eager=enforce_eager,
+                enable_prompt_embeds=enable_prompt_embeds,
+                gpu_memory_utilization=0.7,
+                async_scheduling=async_scheduling,
+                distributed_executor_backend=model_executor,
         ) as vllm_model:
             if enable_prompt_embeds:
-                vllm_outputs = vllm_model.generate_greedy(prompt_embeds, max_tokens)
+                vllm_outputs = vllm_model.generate_greedy(
+                    prompt_embeds, max_tokens)
                 vllm_outputs = _fix_prompt_embed_outputs(
-                    vllm_outputs, hf_model, example_prompts
-                )
+                    vllm_outputs, hf_model, example_prompts)
             else:
-                vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
+                vllm_outputs = vllm_model.generate_greedy(
+                    example_prompts, max_tokens)
 
         check_outputs_equal(
             outputs_0_lst=hf_outputs,
@@ -123,18 +117,21 @@ def test_models(
 
 @multi_gpu_test(num_gpus=2)
 @pytest.mark.parametrize(
-    "model, distributed_executor_backend, attention_backend, test_suite, extra_env",
-    [
+    "model, distributed_executor_backend, attention_backend, "
+    "test_suite, extra_env", [
         ("distilbert/distilgpt2", "ray", "", "L4", {}),
         ("distilbert/distilgpt2", "mp", "", "L4", {}),
-        ("distilbert/distilgpt2", "ray", "", "L4", {"VLLM_SLEEP_WHEN_IDLE": "1"}),
-        ("distilbert/distilgpt2", "mp", "", "L4", {"VLLM_SLEEP_WHEN_IDLE": "1"}),
+        ("distilbert/distilgpt2", "ray", "", "L4", {
+            "VLLM_SLEEP_WHEN_IDLE": "1"
+        }),
+        ("distilbert/distilgpt2", "mp", "", "L4", {
+            "VLLM_SLEEP_WHEN_IDLE": "1"
+        }),
         ("meta-llama/Llama-3.2-1B-Instruct", "ray", "", "L4", {}),
         ("meta-llama/Llama-3.2-1B-Instruct", "mp", "", "L4", {}),
         ("distilbert/distilgpt2", "ray", "", "A100", {}),
         ("distilbert/distilgpt2", "mp", "", "A100", {}),
-    ],
-)
+    ])
 @pytest.mark.parametrize("enable_prompt_embeds", [True, False])
 def test_models_distributed(
     monkeypatch: pytest.MonkeyPatch,
@@ -152,14 +149,11 @@ def test_models_distributed(
         pytest.skip(f"Skip test for {test_suite}")
 
     with monkeypatch.context() as monkeypatch_context:
-        if (
-            model == "meta-llama/Llama-3.2-1B-Instruct"
-            and distributed_executor_backend == "ray"
-            and attention_backend == ""
-            and test_suite == "L4"
-        ):  # noqa
+        if model == "meta-llama/Llama-3.2-1B-Instruct" and distributed_executor_backend == "ray" and attention_backend == "" and test_suite == "L4":  # noqa
             if enable_prompt_embeds:
-                pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")
+                pytest.skip(
+                    "enable_prompt_embeds does not work with ray compiled dag."
+                )
             monkeypatch_context.setenv("VLLM_USE_RAY_SPMD_WORKER", "1")
             monkeypatch_context.setenv("VLLM_USE_RAY_COMPILED_DAG", "1")
 
@@ -181,26 +175,30 @@ def test_models_distributed(
         # will hurt multiprocessing backend with fork method
         # (the default method).
         with vllm_runner(
-            model,
-            dtype=dtype,
-            tensor_parallel_size=2,
-            distributed_executor_backend=distributed_executor_backend,
-            enable_prompt_embeds=enable_prompt_embeds,
-            gpu_memory_utilization=0.7,
+                model,
+                dtype=dtype,
+                tensor_parallel_size=2,
+                distributed_executor_backend=distributed_executor_backend,
+                enable_prompt_embeds=enable_prompt_embeds,
+                gpu_memory_utilization=0.7,
         ) as vllm_model:
             if enable_prompt_embeds:
                 with hf_runner(model, dtype=dtype) as hf_model:
                     with torch.no_grad():
-                        prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
-                    vllm_outputs = vllm_model.generate_greedy(prompt_embeds, max_tokens)
+                        prompt_embeds = hf_model.get_prompt_embeddings(
+                            example_prompts)
+                    vllm_outputs = vllm_model.generate_greedy(
+                        prompt_embeds, max_tokens)
                     vllm_outputs = _fix_prompt_embed_outputs(
-                        vllm_outputs, hf_model, example_prompts
-                    )
-                    hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
+                        vllm_outputs, hf_model, example_prompts)
+                    hf_outputs = hf_model.generate_greedy(
+                        example_prompts, max_tokens)
             else:
-                vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
+                vllm_outputs = vllm_model.generate_greedy(
+                    example_prompts, max_tokens)
                 with hf_runner(model, dtype=dtype) as hf_model:
-                    hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
+                    hf_outputs = hf_model.generate_greedy(
+                        example_prompts, max_tokens)
 
     check_outputs_equal(
         outputs_0_lst=hf_outputs,
@@ -211,23 +209,27 @@ def test_models_distributed(
 
 
 def test_failed_model_execution(vllm_runner, monkeypatch) -> None:
+
     from vllm.envs import VLLM_USE_V1
 
     if not VLLM_USE_V1:
         pytest.skip("Skipping V0 test, dump input not supported")
 
     # Needed to mock an error in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
 
-    with vllm_runner("facebook/opt-125m", enforce_eager=True) as vllm_model:
+    with vllm_runner('facebook/opt-125m', enforce_eager=True) as vllm_model:
         if isinstance(vllm_model.llm.llm_engine, LLMEngineV1):
             v1_test_failed_model_execution(vllm_model)
 
 
 def v1_test_failed_model_execution(vllm_model):
+
     engine = vllm_model.llm.llm_engine
-    mocked_execute_model = Mock(side_effect=RuntimeError("Mocked Critical Error"))
-    engine.engine_core.engine_core.model_executor.execute_model = mocked_execute_model
+    mocked_execute_model = Mock(
+        side_effect=RuntimeError("Mocked Critical Error"))
+    engine.engine_core.engine_core.model_executor.execute_model =\
+                mocked_execute_model
 
     with pytest.raises(RuntimeError) as exc_info:
         prompts = [

tests/basic_correctness/test_cpu_offload.py

@@ -5,6 +5,5 @@ from ..utils import compare_two_settings
 
 
 def test_cpu_offload():
-    compare_two_settings(
-        "meta-llama/Llama-3.2-1B-Instruct", [], ["--cpu-offload-gb", "1"]
-    )
+    compare_two_settings("meta-llama/Llama-3.2-1B-Instruct", [],
+                         ["--cpu-offload-gb", "1"])

tests/basic_correctness/test_cumem.py

@@ -23,13 +23,13 @@ def test_python_error():
     tensors = []
     with allocator.use_memory_pool():
         # allocate 70% of the total memory
-        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device='cuda')
         tensors.append(x)
     # release the memory
     allocator.sleep()
 
     # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device='cuda')
     tensors.append(y)
     with pytest.raises(RuntimeError):
         # when the allocator is woken up, it should raise an error
@@ -41,17 +41,17 @@ def test_python_error():
 def test_basic_cumem():
     # some tensors from default memory pool
     shape = (1024, 1024)
-    x = torch.empty(shape, device="cuda")
+    x = torch.empty(shape, device='cuda')
     x.zero_()
 
     # some tensors from custom memory pool
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
         # custom memory pool
-        y = torch.empty(shape, device="cuda")
+        y = torch.empty(shape, device='cuda')
         y.zero_()
         y += 1
-        z = torch.empty(shape, device="cuda")
+        z = torch.empty(shape, device='cuda')
         z.zero_()
         z += 2
 
@@ -74,16 +74,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
-        weight = torch.eye(1024, device="cuda")
+        weight = torch.eye(1024, device='cuda')
     with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device="cuda")
+        cache = torch.empty(1024, 1024, device='cuda')
 
     def model(x):
         out = x @ weight
-        cache[: out.size(0)].copy_(out)
+        cache[:out.size(0)].copy_(out)
         return out + 1
 
-    x = torch.empty(128, 1024, device="cuda")
+    x = torch.empty(128, 1024, device='cuda')
 
     # warmup
     model(x)
@@ -109,7 +109,7 @@ def test_cumem_with_cudagraph():
     model_graph.replay()
 
     # cache content is as expected
-    assert torch.allclose(x, cache[: x.size(0)])
+    assert torch.allclose(x, cache[:x.size(0)])
 
     # output content is as expected
     assert torch.allclose(y, x + 1)
@@ -123,8 +123,7 @@ def test_cumem_with_cudagraph():
         ("meta-llama/Llama-3.2-1B", True),
         # sleep mode with pytorch checkpoint
         ("facebook/opt-125m", True),
-    ],
-)
+    ])
 def test_end_to_end(monkeypatch: pytest.MonkeyPatch, model: str, use_v1: bool):
     with monkeypatch.context() as m:
         assert use_v1

tests/benchmarks/test_latency_cli.py

@@ -10,18 +10,8 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 @pytest.mark.benchmark
 def test_bench_latency():
     command = [
-        "vllm",
-        "bench",
-        "latency",
-        "--model",
-        MODEL_NAME,
-        "--input-len",
-        "32",
-        "--output-len",
-        "1",
-        "--enforce-eager",
-        "--load-format",
-        "dummy",
+        "vllm", "bench", "latency", "--model", MODEL_NAME, "--input-len", "32",
+        "--output-len", "1", "--enforce-eager", "--load-format", "dummy"
     ]
     result = subprocess.run(command, capture_output=True, text=True)
     print(result.stdout)

tests/benchmarks/test_random_dataset.py

@@ -7,11 +7,8 @@ import numpy as np
 import pytest
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
-from vllm.benchmarks.datasets import (
-    RandomDataset,
-    RandomMultiModalDataset,
-    SampleRequest,
-)
+from vllm.benchmarks.datasets import (RandomDataset, RandomMultiModalDataset,
+                                      SampleRequest)
 
 
 @pytest.fixture(scope="session")
@@ -30,9 +27,11 @@ class Params(NamedTuple):
 
 @pytest.fixture(scope="session")
 def random_dataset_params() -> Params:
-    return Params(
-        num_requests=16, prefix_len=7, range_ratio=0.3, input_len=50, output_len=20
-    )
+    return Params(num_requests=16,
+                  prefix_len=7,
+                  range_ratio=0.3,
+                  input_len=50,
+                  output_len=20)
 
 
 def _fingerprint_sample(req: SampleRequest) -> tuple[str, int, int]:
@@ -40,15 +39,13 @@ def _fingerprint_sample(req: SampleRequest) -> tuple[str, int, int]:
     return (req.prompt, req.prompt_len, req.expected_output_len)
 
 
-def _collect_samples(
-    dataset: RandomDataset,
-    tokenizer: PreTrainedTokenizerBase,
-    num_requests: int = 16,
-    prefix_len: int = 7,
-    range_ratio: float = 0.3,
-    input_len: int = 50,
-    output_len: int = 20,
-) -> list[tuple[str, int, int]]:
+def _collect_samples(dataset: RandomDataset,
+                    tokenizer: PreTrainedTokenizerBase,
+                     num_requests: int = 16,
+                     prefix_len: int = 7,
+                     range_ratio: float = 0.3,
+                     input_len: int = 50,
+                     output_len: int = 20) -> list[tuple[str, int, int]]:
     samples = dataset.sample(
         tokenizer=tokenizer,
         num_requests=num_requests,
@@ -62,8 +59,8 @@ def _collect_samples(
 
 @pytest.mark.benchmark
 def test_random_dataset_same_seed(
-    hf_tokenizer: PreTrainedTokenizerBase, random_dataset_params: Params
-) -> None:
+        hf_tokenizer: PreTrainedTokenizerBase,
+        random_dataset_params: Params) -> None:
     """Same seed should yield identical outputs, even if global RNGs change.
 
     This guards against accidental reliance on Python's random or np.random
@@ -73,15 +70,13 @@ def test_random_dataset_same_seed(
     common_seed = 123
     dataset_a = RandomDataset(random_seed=common_seed)
     dataset_b = RandomDataset(random_seed=common_seed)
-    a = _collect_samples(
-        dataset_a,
-        hf_tokenizer,
-        num_requests=p.num_requests,
-        prefix_len=p.prefix_len,
-        range_ratio=p.range_ratio,
-        input_len=p.input_len,
-        output_len=p.output_len,
-    )
+    a = _collect_samples(dataset_a,
+                         hf_tokenizer,
+                         num_requests=p.num_requests,
+                         prefix_len=p.prefix_len,
+                         range_ratio=p.range_ratio,
+                         input_len=p.input_len,
+                         output_len=p.output_len)
 
     # Perturb global RNG state to ensure isolation
     random.seed(999)
@@ -89,50 +84,43 @@ def test_random_dataset_same_seed(
     np.random.seed(888)
     _ = [np.random.random() for _ in range(100)]
 
-    b = _collect_samples(
-        dataset_b,
-        hf_tokenizer,
-        num_requests=p.num_requests,
-        prefix_len=p.prefix_len,
-        range_ratio=p.range_ratio,
-        input_len=p.input_len,
-        output_len=p.output_len,
-    )
+    b = _collect_samples(dataset_b,
+                         hf_tokenizer,
+                         num_requests=p.num_requests,
+                         prefix_len=p.prefix_len,
+                         range_ratio=p.range_ratio,
+                         input_len=p.input_len,
+                         output_len=p.output_len)
     assert a == b
 
-
 @pytest.mark.benchmark
 def test_random_dataset_different_seeds(
-    hf_tokenizer: PreTrainedTokenizerBase, random_dataset_params: Params
-) -> None:
+        hf_tokenizer: PreTrainedTokenizerBase,
+        random_dataset_params: Params) -> None:
     """Different seeds should change outputs with overwhelming likelihood."""
     p = random_dataset_params
     seed_a = 0
     dataset_a = RandomDataset(random_seed=seed_a)
-    a = _collect_samples(
-        dataset_a,
-        hf_tokenizer,
-        num_requests=p.num_requests,
-        prefix_len=p.prefix_len,
-        range_ratio=p.range_ratio,
-        input_len=p.input_len,
-        output_len=p.output_len,
-    )
+    a = _collect_samples(dataset_a,
+                         hf_tokenizer,
+                         num_requests=p.num_requests,
+                         prefix_len=p.prefix_len,
+                         range_ratio=p.range_ratio,
+                         input_len=p.input_len,
+                         output_len=p.output_len)
 
     seed_b = 999
     dataset_b = RandomDataset(random_seed=seed_b)
     # Perturb global RNG with same seed as dataset_a to ensure isolation
     random.seed(seed_a)
     np.random.seed(seed_a)
-    b = _collect_samples(
-        dataset_b,
-        hf_tokenizer,
-        num_requests=p.num_requests,
-        prefix_len=p.prefix_len,
-        range_ratio=p.range_ratio,
-        input_len=p.input_len,
-        output_len=p.output_len,
-    )
+    b = _collect_samples(dataset_b,
+                         hf_tokenizer,
+                         num_requests=p.num_requests,
+                         prefix_len=p.prefix_len,
+                         range_ratio=p.range_ratio,
+                         input_len=p.input_len,
+                         output_len=p.output_len)
     assert a != b
 
 
@@ -140,7 +128,6 @@ def test_random_dataset_different_seeds(
 # RandomMultiModalDataset tests
 # -----------------------------
 
-
 def _mm_fingerprint_sample(
     req: SampleRequest,
 ) -> tuple[str, int, int, int, list[str]]:
@@ -165,13 +152,8 @@ def _mm_fingerprint_sample(
             item_prefixes.append(f"video:{url[:22]}")
         else:
             item_prefixes.append("unknown:")
-    return (
-        req.prompt,
-        req.prompt_len,
-        req.expected_output_len,
-        len(items),
-        item_prefixes,
-    )
+    return (req.prompt, req.prompt_len, req.expected_output_len, len(items),
+            item_prefixes)
 
 
 def _collect_mm_samples(
@@ -232,7 +214,6 @@ def test_random_mm_different_seeds(
     fb = [_mm_fingerprint_sample(s) for s in b]
     assert fa != fb
 
-
 @pytest.mark.benchmark
 def test_random_mm_respects_limits(
     hf_tokenizer: PreTrainedTokenizerBase,
@@ -290,9 +271,9 @@ def test_random_mm_zero_items(hf_tokenizer: PreTrainedTokenizerBase) -> None:
     for s in samples:
         assert s.multi_modal_data == []
 
-
 @pytest.mark.benchmark
-def test_random_mm_num_items_per_prompt(hf_tokenizer: PreTrainedTokenizerBase) -> None:
+def test_random_mm_num_items_per_prompt(
+    hf_tokenizer: PreTrainedTokenizerBase) -> None:
     ds = RandomMultiModalDataset(random_seed=0)
     # Fixed number of images per prompt
     # set num_mm_items_range_ratio to 0.0
@@ -319,6 +300,7 @@ def test_random_mm_num_items_per_prompt(hf_tokenizer: PreTrainedTokenizerBase) -
 def test_random_mm_bucket_config_not_mutated(
     hf_tokenizer: PreTrainedTokenizerBase,
 ) -> None:
+
     ds = RandomMultiModalDataset(random_seed=0)
     # This bucket config is not normalized to sum to 1
     # and has more buckets than requested images
@@ -339,6 +321,7 @@ def test_random_mm_bucket_config_not_mutated(
     # Ensure the original dict content is unchanged
     assert original == snapshot
 
+
     # Vary number of mm items per prompt
     # set num_mm_items_range_ratio to 0.5
     samples_varying_items = _collect_mm_samples(

tests/benchmarks/test_serve_cli.py

@@ -11,7 +11,9 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 @pytest.fixture(scope="module")
 def server():
-    args = ["--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"]
+    args = [
+        "--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"
+    ]
 
     with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
@@ -44,7 +46,6 @@ def test_bench_serve(server):
 
     assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
 
-
 @pytest.mark.benchmark
 def test_bench_serve_chat(server):
     command = [

tests/benchmarks/test_throughput_cli.py

@@ -10,18 +10,8 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 @pytest.mark.benchmark
 def test_bench_throughput():
     command = [
-        "vllm",
-        "bench",
-        "throughput",
-        "--model",
-        MODEL_NAME,
-        "--input-len",
-        "32",
-        "--output-len",
-        "1",
-        "--enforce-eager",
-        "--load-format",
-        "dummy",
+        "vllm", "bench", "throughput", "--model", MODEL_NAME, "--input-len",
+        "32", "--output-len", "1", "--enforce-eager", "--load-format", "dummy"
     ]
     result = subprocess.run(command, capture_output=True, text=True)
     print(result.stdout)

tests/compile/backend.py

@@ -23,7 +23,8 @@ class LazyInitPass(InductorPass):
     and then immediately invoke it.
     """
 
-    def __init__(self, pass_cls: type[VllmInductorPass], vllm_config: VllmConfig):
+    def __init__(self, pass_cls: type[VllmInductorPass],
+                 vllm_config: VllmConfig):
         self.pass_cls = pass_cls
         self.vllm_config = weakref.proxy(vllm_config)  # avoid cycle
 
@@ -44,18 +45,20 @@ class TestBackend:
     Inductor config is default-initialized from VllmConfig.CompilationConfig.
     """
 
-    def __init__(self, *passes: Union[InductorPass, Callable[[fx.Graph], None]]):
+    def __init__(self, *passes: Union[InductorPass, Callable[[fx.Graph],
+                                                             None]]):
         self.custom_passes = list(passes)
         compile_config = get_current_vllm_config().compilation_config
         self.inductor_config = compile_config.inductor_compile_config
-        self.inductor_config["force_disable_caches"] = True
-        self.inductor_config["post_grad_custom_post_pass"] = self.post_pass
+        self.inductor_config['force_disable_caches'] = True
+        self.inductor_config['post_grad_custom_post_pass'] = self.post_pass
 
     def __call__(self, graph: fx.GraphModule, example_inputs):
         self.graph_pre_compile = deepcopy(graph)
         from torch._inductor.compile_fx import compile_fx
-
-        return compile_fx(graph, example_inputs, config_patches=self.inductor_config)
+        return compile_fx(graph,
+                          example_inputs,
+                          config_patches=self.inductor_config)
 
     @with_pattern_match_debug
     def post_pass(self, graph: fx.Graph):
@@ -79,7 +82,8 @@ class TestBackend:
             assert num_pre > 0, f"Op {op.name()} not found in pre-pass graph"
             assert num_pre > num_post, f"All nodes remain for op {op.name()}"
             if fully_replaced:
-                assert num_post == 0, f"Unexpected op {op.name()} in post-pass graph"
+                assert num_post == 0, \
+                    f"Unexpected op {op.name()} in post-pass graph"
 
     def check_after_ops(self, ops: Sequence[OpOverload]):
         for op in ops:

tests/compile/piecewise/test_full_cudagraph.py

@@ -38,8 +38,8 @@ test_params_full_cudagraph = []
 MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]
 for mla_backend in MLA_backends:
     test_params_full_cudagraph.append(
-        pytest.param(("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend]))
-    )
+        pytest.param(
+            ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend])))
 
 # Qwen/Qwen2-1.5B-Instruct with other backends
 other_backend_configs = [
@@ -47,8 +47,7 @@ other_backend_configs = [
 ]
 for backend_config in other_backend_configs:
     test_params_full_cudagraph.append(
-        pytest.param(("Qwen/Qwen2-1.5B-Instruct", backend_config))
-    )
+        pytest.param(("Qwen/Qwen2-1.5B-Instruct", backend_config)))
 
 
 @pytest.fixture(scope="class")
@@ -56,10 +55,8 @@ def llm_pair(request):
     model, backend_config = request.param
 
     # Dynamically skip test if GPU capability is not met
-    if (
-        backend_config.specific_gpu_arch
-        and backend_config.specific_gpu_arch != current_platform.get_device_capability()
-    ):
+    if backend_config.specific_gpu_arch and backend_config.specific_gpu_arch\
+        != current_platform.get_device_capability():
         if backend_config.specific_gpu_arch == (9, 0):
             pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
         elif backend_config.specific_gpu_arch == (10, 0):
@@ -79,7 +76,8 @@ def llm_pair(request):
             trust_remote_code=True,
             max_model_len=1024,
             max_num_seqs=128,
-            compilation_config=CompilationConfig(**backend_config.comp_config),
+            compilation_config=\
+                CompilationConfig(**backend_config.comp_config),
             generation_config="vllm",
             seed=42,
         )
@@ -115,22 +113,20 @@ class TestFullCUDAGraph:
     meaning there would be multiple LLM instances hogging memory simultaneously.
     """
 
-    @pytest.mark.parametrize(
-        ("batch_size", "max_tokens"),
-        [
-            (1, 10),
-            (7, 10),
-            (16, 10),
-            (25, 10),
-            (32, 10),
-            (45, 10),
-            (64, 10),
-            (123, 10),
-            (8, 5),
-            (8, 30),
-        ],
-    )
-    def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]):
+    @pytest.mark.parametrize(("batch_size", "max_tokens"), [
+        (1, 10),
+        (7, 10),
+        (16, 10),
+        (25, 10),
+        (32, 10),
+        (45, 10),
+        (64, 10),
+        (123, 10),
+        (8, 5),
+        (8, 30),
+    ])
+    def test_full_cudagraph(self, batch_size, max_tokens,
+                            llm_pair: tuple[LLM, LLM]):
         """
         Test various batch sizes and max_tokens to ensure that the
         full cudagraph compilation works for padded cases too.
@@ -141,34 +137,26 @@ class TestFullCUDAGraph:
         prompts = ["the quick brown fox"] * batch_size
         # Use purely greedy decoding to avoid top-p truncation sensitivity
         # that can amplify tiny numeric differences across runtimes.
-        sampling_params = SamplingParams(
-            temperature=0.0, max_tokens=max_tokens, top_p=1.0
-        )
+        sampling_params = SamplingParams(temperature=0.0,
+                                         max_tokens=max_tokens,
+                                         top_p=1.0)
 
         piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
         full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
 
         # Check that all responses are the same
-        for piecewise_res, full_res in zip(piecewise_responses, full_responses):
-            assert (
-                piecewise_res.outputs[0].text.lower()
-                == full_res.outputs[0].text.lower()
-            )
+        for piecewise_res, full_res in zip(piecewise_responses,
+                                           full_responses):
+            assert piecewise_res.outputs[0].text.lower() == \
+                full_res.outputs[0].text.lower()
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
 def test_full_cudagraph_with_invalid_backend():
-    with (
-        temporary_environ(
-            {
-                "VLLM_USE_V1": "1",
-                "VLLM_ATTENTION_BACKEND": "FLEX_ATTENTION",
-                # Flex_Attention is not supported with full cuda graph
-            }
-        ),
-        pytest.raises(RuntimeError),
-    ):
-        LLM(
-            model="Qwen/Qwen2-1.5B-Instruct",
-            compilation_config=CompilationConfig(cudagraph_mode="FULL"),
-        )
+    with temporary_environ({
+            "VLLM_USE_V1": "1",
+            "VLLM_ATTENTION_BACKEND": "FLEX_ATTENTION"
+            # Flex_Attention is not supported with full cuda graph
+    }), pytest.raises(RuntimeError):
+        LLM(model="Qwen/Qwen2-1.5B-Instruct",
+            compilation_config=CompilationConfig(cudagraph_mode="FULL"))

tests/compile/piecewise/test_multiple_graphs.py

@@ -10,14 +10,10 @@ from torch import nn
 
 from vllm.compilation.backends import set_model_tag
 from vllm.compilation.counter import compilation_counter
-from vllm.compilation.decorators import ignore_torch_compile, support_torch_compile
-from vllm.config import (
-    CompilationConfig,
-    CompilationLevel,
-    CUDAGraphMode,
-    VllmConfig,
-    set_current_vllm_config,
-)
+from vllm.compilation.decorators import (ignore_torch_compile,
+                                         support_torch_compile)
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
+                         VllmConfig, set_current_vllm_config)
 from vllm.forward_context import BatchDescriptor, set_forward_context
 
 # This import automatically registers `torch.ops.silly.attention`
@@ -31,7 +27,12 @@ RANDOM_SEED = 0
 
 @support_torch_compile
 class ParentModel(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -39,6 +40,7 @@ class ParentModel(nn.Module):
 
 
 class Attention(nn.Module):
+
     def __init__(self, mlp_size: int, hidden_size: int) -> None:
         super().__init__()
         self.pre_attn = nn.Linear(mlp_size, hidden_size, bias=False)
@@ -49,21 +51,17 @@ class Attention(nn.Module):
         nn.init.xavier_normal_(
             self.pre_attn.weight.data,
             generator=torch.Generator().manual_seed(RANDOM_SEED),
-            gain=0.001,
-        )
+            gain=0.001)
         nn.init.xavier_normal_(
             self.post_attn.weight.data,
             generator=torch.Generator().manual_seed(RANDOM_SEED),
-            gain=0.001,
-        )
+            gain=0.001)
 
     def rms_norm_ref(self, x: torch.Tensor) -> torch.Tensor:
         x_f32 = x.float()
-        return (
-            x_f32
-            * torch.rsqrt(torch.mean(x_f32.square(), dim=-1, keepdim=True) + 1e-6)
-            * self.rms_norm_weight
-        ).to(x.dtype)
+        return (x_f32 * torch.rsqrt(
+            torch.mean(x_f32.square(), dim=-1, keepdim=True) + 1e-6) *
+                self.rms_norm_weight).to(x.dtype)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.pre_attn(x)
@@ -78,15 +76,14 @@ class Attention(nn.Module):
 
 @support_torch_compile
 class CompiledAttention(nn.Module):
-    def __init__(
-        self,
-        *,
-        mlp_size: int,
-        hidden_size: int,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        **kwargs,
-    ) -> None:
+
+    def __init__(self,
+                 *,
+                 mlp_size: int,
+                 hidden_size: int,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
         self.attn = Attention(mlp_size, hidden_size)
 
@@ -96,21 +93,21 @@ class CompiledAttention(nn.Module):
 
 @support_torch_compile
 class CompiledAttentionTwo(CompiledAttention):
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.attn(x) + x
 
 
 @ignore_torch_compile
 class SimpleModelWithTwoGraphs(ParentModel):
-    def __init__(
-        self,
-        *,
-        mlp_size: int,
-        hidden_size: int,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        **kwargs,
-    ) -> None:
+
+    def __init__(self,
+                 *,
+                 mlp_size: int,
+                 hidden_size: int,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         # Test will fail without set_model_tag here with error:
         # "ValueError: too many values to unpack (expected 3)"
@@ -145,45 +142,32 @@ class SimpleModelWithTwoGraphs(ParentModel):
 
 
 @torch.inference_mode
-def run_model(
-    vllm_config: VllmConfig,
-    model: nn.Module,
-    inputs: torch.Tensor,
-    cudagraph_runtime_mode: CUDAGraphMode,
-):
+def run_model(vllm_config: VllmConfig, model: nn.Module, inputs: torch.Tensor,
+              cudagraph_runtime_mode: CUDAGraphMode):
     with set_forward_context({}, vllm_config=vllm_config):
         # warmup for the model with cudagraph_mode NONE
         model(inputs)
 
         # simulate cudagraphs capturing
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             model(inputs[:2])
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=1,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=1, )):
             model(inputs[:1])
 
         # simulate cudagraphs replay
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             output = model(inputs[:2])
 
         output = output.cpu()
@@ -194,104 +178,82 @@ def test_multi_graph_piecewise_compile_outputs_equal():
     outputs = []
 
     # piecewise compile
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=True,
-            splitting_ops=["silly.attention"],
-            cudagraph_capture_sizes=[1, 2],
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=True,
+        splitting_ops=["silly.attention"],
+        cudagraph_capture_sizes=[1, 2],
+    ))
     cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
 
     with set_current_vllm_config(vllm_config):
-        model = (
-            SimpleModelWithTwoGraphs(
-                mlp_size=MLP_SIZE,
-                hidden_size=HIDDEN_SIZE,
-                vllm_config=vllm_config,
-                prefix="",
-            )
-            .eval()
-            .cuda()
-        )
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
 
     # Pre-allocate memory for CUDAGraph which expects
     # static tensor addresses
     inputs = torch.randn(BATCH_SIZE, MLP_SIZE).cuda()
 
     with compilation_counter.expect(
-        num_graphs_seen=2,  # two graphs for the model
-        num_piecewise_graphs_seen=6,
-        # attn_one, attn_two each has 3 piecewise graphs
-        # (pre attn, post attn, silly_attention) each
-        num_piecewise_capturable_graphs_seen=4,
-        # attn_one, attn_two has pre attn and post attn each, total=4
-        num_backend_compilations=4,  # num_piecewise_capturable_graphs_seen
-        num_cudagraph_captured=8,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=2,  # two graphs for the model
+            num_piecewise_graphs_seen=6,
+            # attn_one, attn_two each has 3 piecewise graphs
+            # (pre attn, post attn, silly_attention) each
+            num_piecewise_capturable_graphs_seen=4,
+            # attn_one, attn_two has pre attn and post attn each, total=4
+            num_backend_compilations=4,  # num_piecewise_capturable_graphs_seen
+            num_cudagraph_captured=8,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
-        outputs.append(run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
+        outputs.append(
+            run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
 
     # no compile or cudagraph
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.NO_COMPILATION,
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.NO_COMPILATION, ))
     cudagraph_runtime_mode = CUDAGraphMode.NONE
 
     with set_current_vllm_config(vllm_config):
-        model = (
-            SimpleModelWithTwoGraphs(
-                mlp_size=MLP_SIZE,
-                hidden_size=HIDDEN_SIZE,
-                vllm_config=vllm_config,
-                prefix="",
-            )
-            .eval()
-            .cuda()
-        )
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
 
     with compilation_counter.expect(
-        num_graphs_seen=0,
-        num_piecewise_graphs_seen=0,
-        num_piecewise_capturable_graphs_seen=0,
-        num_backend_compilations=0,
-        num_cudagraph_captured=0,
+            num_graphs_seen=0,
+            num_piecewise_graphs_seen=0,
+            num_piecewise_capturable_graphs_seen=0,
+            num_backend_compilations=0,
+            num_cudagraph_captured=0,
     ):
-        outputs.append(run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
+        outputs.append(
+            run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
 
     # piecewise compile without CUDA graph
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=False,
-            splitting_ops=["silly.attention"],
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=False,
+        splitting_ops=["silly.attention"],
+    ))
     cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
 
     with set_current_vllm_config(vllm_config):
-        model = (
-            SimpleModelWithTwoGraphs(
-                mlp_size=MLP_SIZE,
-                hidden_size=HIDDEN_SIZE,
-                vllm_config=vllm_config,
-                prefix="",
-            )
-            .eval()
-            .cuda()
-        )
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
 
     with compilation_counter.expect(
-        num_graphs_seen=2,
-        num_piecewise_graphs_seen=6,
-        num_piecewise_capturable_graphs_seen=4,
-        num_backend_compilations=4,
-        num_cudagraph_captured=0,  # no cudagraph captured
+            num_graphs_seen=2,
+            num_piecewise_graphs_seen=6,
+            num_piecewise_capturable_graphs_seen=4,
+            num_backend_compilations=4,
+            num_cudagraph_captured=0,  # no cudagraph captured
     ):
-        outputs.append(run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
+        outputs.append(
+            run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
 
     # Generally don't expect outputs with and without inductor
     # to be bitwise equivalent

tests/compile/piecewise/test_simple.py

@@ -11,13 +11,8 @@ from torch import nn
 
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import (
-    CompilationConfig,
-    CompilationLevel,
-    CUDAGraphMode,
-    VllmConfig,
-    set_current_vllm_config,
-)
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
+                         VllmConfig, set_current_vllm_config)
 from vllm.envs import VLLM_USE_V1
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.utils import is_torch_equal_or_newer
@@ -28,7 +23,12 @@ from ..silly_attention import get_global_counter, reset_global_counter
 
 @support_torch_compile
 class SillyModel(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -60,65 +60,53 @@ def _run_simple_model(
     expected_num_backend_compilations,
     expected_num_cudagraph_captured,
 ):
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=True,
-            use_inductor=use_inductor,
-            splitting_ops=splitting_ops,
-            use_inductor_graph_partition=use_inductor_graph_partition,
-            cudagraph_copy_inputs=True,
-            cudagraph_capture_sizes=[1, 2],
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=True,
+        use_inductor=use_inductor,
+        splitting_ops=splitting_ops,
+        use_inductor_graph_partition=use_inductor_graph_partition,
+        cudagraph_copy_inputs=True,
+        cudagraph_capture_sizes=[1, 2],
+    ))
     with set_current_vllm_config(vllm_config):
-        model = SillyModel(vllm_config=vllm_config, prefix="")
+        model = SillyModel(vllm_config=vllm_config, prefix='')
 
     inputs = torch.randn(100).cuda()
 
-    with (
-        compilation_counter.expect(
+    with compilation_counter.expect(
             num_graphs_seen=1,  # one graph for the model
             num_piecewise_graphs_seen=expected_num_piecewise_graphs_seen,
-            num_piecewise_capturable_graphs_seen=expected_num_piecewise_capturable_graphs_seen,
+            num_piecewise_capturable_graphs_seen=
+            expected_num_piecewise_capturable_graphs_seen,
             num_backend_compilations=expected_num_backend_compilations,
             num_cudagraph_captured=expected_num_cudagraph_captured,
-        ),
-        set_forward_context(None, vllm_config=vllm_config),
-    ):  # background context
+    ), set_forward_context(None,
+                           vllm_config=vllm_config):  # background context
         # warm up with background context
         model(inputs)
 
         # capturing/replaying should under context of cudagraph dispatching
         with set_forward_context(
-            None,
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+                None,
+                vllm_config=vllm_config,
+                cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+                batch_descriptor=BatchDescriptor(num_tokens=2, )):
             model(torch.randn(2).cuda())
         with set_forward_context(
-            None,
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=1,
-            ),
-        ):
+                None,
+                vllm_config=vllm_config,
+                cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+                batch_descriptor=BatchDescriptor(num_tokens=1, )):
             model(torch.randn(1).cuda())
 
         input = torch.zeros(2).cuda()
         reset_global_counter()
         with set_forward_context(
-            None,
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+                None,
+                vllm_config=vllm_config,
+                cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+                batch_descriptor=BatchDescriptor(num_tokens=2, )):
             output = model(input)
         assert get_global_counter() == 2
         assert torch.allclose(output.cpu(), torch.tensor([19.0, 19.0]))
@@ -132,14 +120,12 @@ def test_simple_piecewise_compile(use_inductor):
         splitting_ops=["silly.attention"],
         use_inductor_graph_partition=False,
         use_inductor=use_inductor,
-        # 2 * num_layers + 1
-        expected_num_piecewise_graphs_seen=5,
-        # 1 + num_layers
-        expected_num_piecewise_capturable_graphs_seen=3,
-        # num_piecewise_capturable_graphs_seen
-        expected_num_backend_compilations=3,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
-        expected_num_cudagraph_captured=6,
+        expected_num_piecewise_graphs_seen=5,  # 2 * num_layers + 1
+        expected_num_piecewise_capturable_graphs_seen=3,  # 1 + num_layers
+        expected_num_backend_compilations=
+        3,  # num_piecewise_capturable_graphs_seen
+        expected_num_cudagraph_captured=
+        6,  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     )
 
 
@@ -148,19 +134,22 @@ def test_simple_piecewise_compile(use_inductor):
 def test_simple_inductor_graph_partition(splitting_ops):
     assert VLLM_USE_V1
     if not is_torch_equal_or_newer("2.9.0.dev"):
-        pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
+        pytest.skip("inductor graph partition is only available "
+                    "in PyTorch 2.9+")
 
     _run_simple_model(
-        # Inductor graph partition automatically resets splitting_ops to an empty list
+        # inductor graph partition automatically resets splitting_ops
+        # to be an empty list
         splitting_ops=splitting_ops,
         use_inductor_graph_partition=True,
         use_inductor=True,
-        # Since not splitting at fx graph level
-        expected_num_piecewise_graphs_seen=1,
-        # Since not splitting at fx graph level
-        expected_num_piecewise_capturable_graphs_seen=1,
-        # Since not splitting at fx graph level
-        expected_num_backend_compilations=1,
-        # Inductor graph partition still captures 6 graph, same as fx graph partition
-        expected_num_cudagraph_captured=6,
+        expected_num_piecewise_graphs_seen=
+        1,  # since not splitting at fx graph level
+        expected_num_piecewise_capturable_graphs_seen=
+        1,  # since not splitting at fx graph level
+        expected_num_backend_compilations=
+        1,  # since not splitting at fx graph level
+        expected_num_cudagraph_captured=
+        6,  # inductor graph partition still captures 6
+        # graph, same as fx graph partition.
     )

tests/compile/piecewise/test_toy_llama.py

@@ -8,7 +8,6 @@ This is a tractable model, the weights and computation are specially designed
 if the config `tractable_init` is set to True. Otherwise, the weights are
 initialized randomly with a fixed seed.
 """
-
 from dataclasses import dataclass
 from typing import Any, Optional
 
@@ -18,13 +17,8 @@ from torch import nn
 
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import (
-    CompilationConfig,
-    CompilationLevel,
-    CUDAGraphMode,
-    VllmConfig,
-    set_current_vllm_config,
-)
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
+                         VllmConfig, set_current_vllm_config)
 from vllm.forward_context import BatchDescriptor, set_forward_context
 
 # This import automatically registers `torch.ops.silly.attention`
@@ -49,14 +43,15 @@ class LlamaConfig:
             factors.append((k, v))
         factors.sort()
         import hashlib
-
-        return hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
+        return hashlib.md5(str(factors).encode(),
+                           usedforsecurity=False).hexdigest()
 
     def __post_init__(self):
         assert self.mlp_size >= self.hidden_size
 
 
 class LlamaMLP(nn.Module):
+
     def __init__(self, config: LlamaConfig) -> None:
         super().__init__()
         self.gate_up_projection = nn.Linear(
@@ -71,31 +66,31 @@ class LlamaMLP(nn.Module):
         )
 
         if config.tractable_init:
-            nn.init.eye_(self.gate_up_projection.weight.data[: config.mlp_size])
-            nn.init.eye_(self.gate_up_projection.weight.data[config.mlp_size :])
+            nn.init.eye_(self.gate_up_projection.weight.data[:config.mlp_size])
+            nn.init.eye_(self.gate_up_projection.weight.data[config.mlp_size:])
             nn.init.eye_(self.down_projection.weight.data)
         else:
-            nn.init.xavier_normal_(
-                self.gate_up_projection.weight.data,
-                generator=torch.Generator().manual_seed(config.random_seed),
-                gain=0.001,
-            )
-            nn.init.xavier_normal_(
-                self.down_projection.weight.data,
-                generator=torch.Generator().manual_seed(config.random_seed),
-                gain=0.001,
-            )
+            nn.init.xavier_normal_(self.gate_up_projection.weight.data,
+                                   generator=torch.Generator().manual_seed(
+                                       config.random_seed),
+                                   gain=0.001)
+            nn.init.xavier_normal_(self.down_projection.weight.data,
+                                   generator=torch.Generator().manual_seed(
+                                       config.random_seed),
+                                   gain=0.001)
 
     def forward(self, x):
         # for tractable_init and positive input, this is
         # essentially an elementwise-square
         x = self.gate_up_projection(x)
-        x = x[:, : x.size(1) // 2] * torch.nn.functional.relu(x[:, x.size(1) // 2 :])
+        x = x[:, :x.size(1) // 2] * torch.nn.functional.relu(
+            x[:, x.size(1) // 2:])
         x = self.down_projection(x)
         return x
 
 
 class LlamaAttention(nn.Module):
+
     def __init__(self, config: LlamaConfig) -> None:
         super().__init__()
         self.qkv_projection = nn.Linear(
@@ -111,25 +106,21 @@ class LlamaAttention(nn.Module):
         )
 
         if config.tractable_init:
-            nn.init.eye_(self.qkv_projection.weight.data[: config.hidden_size])
-            nn.init.eye_(
-                self.qkv_projection.weight.data[
-                    config.hidden_size : 2 * config.hidden_size
-                ]
-            )
-            nn.init.eye_(self.qkv_projection.weight.data[2 * config.hidden_size :])
+            nn.init.eye_(self.qkv_projection.weight.data[:config.hidden_size])
+            nn.init.eye_(self.qkv_projection.weight.data[config.hidden_size:2 *
+                                                         config.hidden_size])
+            nn.init.eye_(self.qkv_projection.weight.data[2 *
+                                                         config.hidden_size:])
             nn.init.eye_(self.output_projection.weight.data)
         else:
-            nn.init.xavier_normal_(
-                self.qkv_projection.weight.data,
-                generator=torch.Generator().manual_seed(config.random_seed),
-                gain=0.001,
-            )
-            nn.init.xavier_normal_(
-                self.output_projection.weight.data,
-                generator=torch.Generator().manual_seed(config.random_seed),
-                gain=0.001,
-            )
+            nn.init.xavier_normal_(self.qkv_projection.weight.data,
+                                   generator=torch.Generator().manual_seed(
+                                       config.random_seed),
+                                   gain=0.001)
+            nn.init.xavier_normal_(self.output_projection.weight.data,
+                                   generator=torch.Generator().manual_seed(
+                                       config.random_seed),
+                                   gain=0.001)
 
     def forward(
         self,
@@ -153,6 +144,7 @@ class LlamaAttention(nn.Module):
 
 
 class LlamaDecoderLayer(nn.Module):
+
     def __init__(self, config: LlamaConfig) -> None:
         super().__init__()
         self.self_attention = LlamaAttention(config)
@@ -172,7 +164,7 @@ class LlamaDecoderLayer(nn.Module):
         - if residual is not None, the outputs are:
             - residual = (hidden_states + residual + 1) * 3 + positions * 2 + hidden_states + residual = (hidden_states + residual) * 4 + positions * 2 + 3
             - hidden_states = (residual + 1) ** 2
-        """  # noqa
+        """ # noqa
         if residual is None:
             residual = hidden_states
             hidden_states = hidden_states + 1
@@ -181,9 +173,8 @@ class LlamaDecoderLayer(nn.Module):
             residual = hidden_states
             hidden_states = hidden_states + 1
 
-        hidden_states = self.self_attention(
-            positions=positions, hidden_states=hidden_states
-        )
+        hidden_states = self.self_attention(positions=positions,
+                                            hidden_states=hidden_states)
 
         hidden_states = hidden_states + residual
         residual = hidden_states
@@ -195,22 +186,20 @@ class LlamaDecoderLayer(nn.Module):
 
 @support_torch_compile
 class LlamaModel(nn.Module):
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        config: LlamaConfig,
-        prefix: str = "",
-        **kwargs,
-    ) -> None:
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 config: LlamaConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
         self.embedding_tokens = nn.Embedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
         )
         self.layers = nn.ModuleList(
-            [LlamaDecoderLayer(config) for _ in range(config.num_layers)]
-        )
+            [LlamaDecoderLayer(config) for _ in range(config.num_layers)])
 
         # this is the initial value of the hidden states
         self.embedding_tokens.weight.data.fill_(config.init_value)
@@ -227,39 +216,34 @@ class LlamaModel(nn.Module):
         return hidden_states
 
 
-def tractable_computation(
-    input_ids: torch.Tensor,
-    positions: torch.Tensor,
-    config: LlamaConfig,
-    init_value: float = 1.0,
-) -> torch.Tensor:
-    hidden_states = (
-        torch.ones(
-            input_ids.size(0),
-            config.hidden_size,
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        * init_value
-    )
+def tractable_computation(input_ids: torch.Tensor,
+                          positions: torch.Tensor,
+                          config: LlamaConfig,
+                          init_value: float = 1.0) -> torch.Tensor:
+    hidden_states = torch.ones(input_ids.size(0),
+                               config.hidden_size,
+                               device=input_ids.device,
+                               dtype=input_ids.dtype) * init_value
 
     # first layer
     residual = hidden_states * 4 + positions.unsqueeze(1) * 2 + 3
-    hidden_states = (residual + 1) ** 2
+    hidden_states = (residual + 1)**2
 
     # following layers
     for _ in range(config.num_layers - 1):
         hidden_states = hidden_states + residual
         residual = hidden_states * 4 + positions.unsqueeze(1) * 2 + 3
-        hidden_states = (residual + 1) ** 2
+        hidden_states = (residual + 1)**2
 
     return hidden_states
 
 
 @torch.inference_mode
-def run_model(
-    llama_config, use_compile: bool, use_inductor: bool, split_attn: bool = False
-) -> torch.Tensor:
+def run_model(llama_config,
+              use_compile: bool,
+              use_inductor: bool,
+              split_attn: bool = False) -> torch.Tensor:
+
     if use_compile:
         compilation_config = CompilationConfig(
             level=CompilationLevel.PIECEWISE,
@@ -272,66 +256,54 @@ def run_model(
         cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
     else:
         compilation_config = CompilationConfig(
-            level=CompilationLevel.NO_COMPILATION,
-        )
+            level=CompilationLevel.NO_COMPILATION, )
         cudagraph_runtime_mode = CUDAGraphMode.NONE
 
-    vllm_config = VllmConfig(
-        compilation_config=compilation_config, additional_config=llama_config
-    )
+    vllm_config = VllmConfig(compilation_config=compilation_config,
+                             additional_config=llama_config)
     with set_current_vllm_config(vllm_config):
-        model = (
-            LlamaModel(config=llama_config, vllm_config=vllm_config, prefix="")
-            .eval()
-            .cuda()
-        )
+        model = LlamaModel(config=llama_config,
+                           vllm_config=vllm_config,
+                           prefix="").eval().cuda()
 
-    with set_forward_context({}, vllm_config=vllm_config):  # background context
+    with set_forward_context({},
+                             vllm_config=vllm_config):  # background context
         B = 16  # max batch size
-        input_ids = torch.randint(0, llama_config.vocab_size, (B,)).cuda()
+        input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
         positions = torch.arange(B).cuda()
 
         # warmup for the model with cudagraph_mode NONE
         model(input_ids, positions)
 
         # simulate cudagraphs capturing
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             model(input_ids[:2], positions[:2])
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=1,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=1, )):
             model(input_ids[:1], positions[:1])
 
         input_ids[:2].zero_()
         # simulate cudagraphs replay
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             output = model(input_ids[:2], positions[:2])
 
         output = output.cpu()
 
         if llama_config.tractable_init:
-            expected_output = tractable_computation(
-                input_ids[:2], positions[:2], llama_config
-            ).cpu()
+            expected_output = tractable_computation(input_ids[:2],
+                                                    positions[:2],
+                                                    llama_config).cpu()
 
             assert torch.allclose(output, expected_output)
         else:
@@ -342,23 +314,27 @@ def run_model(
 def test_toy_llama(use_inductor: bool):
     # compare output with and without piecewise compilation
 
-    llama_config = LlamaConfig(
-        hidden_size=128, mlp_size=256, vocab_size=128, num_layers=12
-    )
+    llama_config = LlamaConfig(hidden_size=128,
+                               mlp_size=256,
+                               vocab_size=128,
+                               num_layers=12)
 
-    tractable_config = LlamaConfig(
-        hidden_size=128, mlp_size=256, vocab_size=128, num_layers=2, tractable_init=True
-    )
+    tractable_config = LlamaConfig(hidden_size=128,
+                                   mlp_size=256,
+                                   vocab_size=128,
+                                   num_layers=2,
+                                   tractable_init=True)
 
     outputs = []
     with compilation_counter.expect(
-        num_graphs_seen=0,
-        num_piecewise_graphs_seen=0,
-        num_piecewise_capturable_graphs_seen=0,
-        num_backend_compilations=0,
-        num_cudagraph_captured=0,
+            num_graphs_seen=0,
+            num_piecewise_graphs_seen=0,
+            num_piecewise_capturable_graphs_seen=0,
+            num_backend_compilations=0,
+            num_cudagraph_captured=0,
     ):
-        outputs.append(run_model(llama_config, use_inductor=False, use_compile=False))
+        outputs.append(
+            run_model(llama_config, use_inductor=False, use_compile=False))
     run_model(tractable_config, use_inductor=False, use_compile=False)
 
     if use_inductor:
@@ -367,44 +343,41 @@ def test_toy_llama(use_inductor: bool):
         kwargs = {"num_eager_compiles": 1, "num_inductor_compiles": 0}
 
     with compilation_counter.expect(
-        # One graph for the model
-        num_graphs_seen=1,
-        num_piecewise_graphs_seen=1,
-        num_piecewise_capturable_graphs_seen=1,
-        # num_piecewise_capturable_graphs_seen
-        num_backend_compilations=1,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
-        num_cudagraph_captured=2,
-        **kwargs,
+            num_graphs_seen=1,  # one graph for the model
+            num_piecewise_graphs_seen=1,
+            num_piecewise_capturable_graphs_seen=1,
+            num_backend_compilations=1,  # num_piecewise_capturable_graphs_seen
+            num_cudagraph_captured=
+            2,  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            **kwargs,
     ):
         outputs.append(
-            run_model(llama_config, use_inductor=use_inductor, use_compile=True)
-        )
+            run_model(llama_config,
+                      use_inductor=use_inductor,
+                      use_compile=True))
     run_model(tractable_config, use_inductor=use_inductor, use_compile=True)
 
     with compilation_counter.expect(
-        num_graphs_seen=1,  # one graph for the model
-        num_piecewise_graphs_seen=2 * llama_config.num_layers + 1,  # 2 * num_layers + 1
-        num_piecewise_capturable_graphs_seen=1
-        + llama_config.num_layers,  # 1 + num_layers
-        num_backend_compilations=1
-        + llama_config.num_layers,  # num_piecewise_capturable_graphs_seen
-        num_cudagraph_captured=2
-        * (
-            1 + llama_config.num_layers
-        ),  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=1,  # one graph for the model
+            num_piecewise_graphs_seen=2 * llama_config.num_layers +
+            1,  # 2 * num_layers + 1
+            num_piecewise_capturable_graphs_seen=1 +
+            llama_config.num_layers,  # 1 + num_layers
+            num_backend_compilations=1 +
+            llama_config.num_layers,  # num_piecewise_capturable_graphs_seen
+            num_cudagraph_captured=2 *
+        (1 + llama_config.num_layers
+         ),  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
         outputs.append(
-            run_model(
-                llama_config,
-                use_inductor=use_inductor,
-                use_compile=True,
-                split_attn=True,
-            )
-        )
-    run_model(
-        tractable_config, use_inductor=use_inductor, use_compile=True, split_attn=True
-    )
+            run_model(llama_config,
+                      use_inductor=use_inductor,
+                      use_compile=True,
+                      split_attn=True))
+    run_model(tractable_config,
+              use_inductor=use_inductor,
+              use_compile=True,
+              split_attn=True)
 
     for i in range(1, len(outputs)):
         assert torch.allclose(outputs[0], outputs[i])
@@ -415,15 +388,17 @@ def benchmark():
     from triton.testing import do_bench
 
     # similar to llama 3.1-8B
-    llama_config = LlamaConfig(
-        hidden_size=4096, mlp_size=14336, vocab_size=128 * 1024, num_layers=32
-    )
+    llama_config = LlamaConfig(hidden_size=4096,
+                               mlp_size=14336,
+                               vocab_size=128 * 1024,
+                               num_layers=32)
 
     # a tiny model to measure the overhead
     # of piecewise cudagraph
-    llama_config = LlamaConfig(
-        hidden_size=40, mlp_size=80, vocab_size=128, num_layers=2
-    )
+    llama_config = LlamaConfig(hidden_size=40,
+                               mlp_size=80,
+                               vocab_size=128,
+                               num_layers=2)
 
     cudagraph_sizes = [1, 2, 4] + [i * 8 for i in range(1, 33)]
 
@@ -449,15 +424,12 @@ def benchmark():
 
         vllm_config = VllmConfig(compilation_config=compilation_config)
         with set_current_vllm_config(vllm_config):
-            model = (
-                LlamaModel(config=llama_config, vllm_config=vllm_config, prefix="")
-                .eval()
-                .cuda()
-                .to(torch.bfloat16)
-            )
+            model = LlamaModel(config=llama_config,
+                               vllm_config=vllm_config,
+                               prefix="").eval().cuda().to(torch.bfloat16)
 
         B = 256  # max batch size
-        input_ids = torch.randint(0, llama_config.vocab_size, (B,)).cuda()
+        input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
         positions = torch.arange(B).cuda().to(torch.bfloat16)
 
         graphs = {}
@@ -479,26 +451,21 @@ def benchmark():
                 # and use it later, because it will look up the name `b` in the
                 # enclosing scope, and the value of `b` will always be 256.
                 # it is fine here, because we only use the lambda function once.
-                runtime = do_bench(
-                    lambda: graphs[b][0](  # noqa
-                        input_ids[:b],  # noqa
-                        positions[:b],  # noqa
-                    )
-                )
+                runtime = do_bench(lambda: graphs[b][0]  # noqa
+                                   (input_ids[:b], positions[:b]))  # noqa
                 piecewise_cudagraph_time[b] = runtime
             else:
                 runtime = do_bench(lambda: graphs[b][0].replay())  # noqa
-                eager_runtime = do_bench(lambda: model(input_ids[:b], positions[:b]))  # noqa
+                eager_runtime = do_bench(
+                    lambda: model(input_ids[:b], positions[:b]))  # noqa
                 full_cudagraph_time[b] = runtime
                 eager_time[b] = eager_runtime
 
     # print in tabular format
     print("batch size\teager mode\tfull cudagraph\tpiecewise cudagraph")
     for b in cudagraph_sizes:
-        print(
-            f"{b}\t{eager_time[b]:.3f}\t{full_cudagraph_time[b]:.3f}"
-            f"\t{piecewise_cudagraph_time[b]:.3f}"
-        )
+        print(f"{b}\t{eager_time[b]:.3f}\t{full_cudagraph_time[b]:.3f}"
+              f"\t{piecewise_cudagraph_time[b]:.3f}")
 
 
 if __name__ == "__main__":

tests/compile/silly_attention.py

@@ -31,9 +31,8 @@ def reset_global_counter():
     _global_counter = 0
 
 
-def silly_attention(
-    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor
-) -> None:
+def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                    out: torch.Tensor) -> None:
     """
     Unified attention implementation that depends on
     all inputs and affects the output.
@@ -48,9 +47,8 @@ def silly_attention(
     out.copy_(q + k + v)
 
 
-def silly_attention_fake(
-    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor
-) -> None:
+def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                         out: torch.Tensor) -> None:
     """Fake implementation for testing"""
     return
 
@@ -62,5 +60,5 @@ direct_register_custom_op(
     mutates_args=["out"],
     fake_impl=silly_attention_fake,
     target_lib=silly_lib,
-    tags=(torch._C.Tag.cudagraph_unsafe,),
+    tags=(torch._C.Tag.cudagraph_unsafe, ),
 )

tests/compile/test_async_tp.py

@@ -8,30 +8,18 @@ import torch
 
 import vllm.envs as envs
 from vllm.compilation.collective_fusion import AsyncTPPass
-from vllm.config import (
-    CompilationConfig,
-    DeviceConfig,
-    ModelConfig,
-    PassConfig,
-    VllmConfig,
-)
-from vllm.distributed import (
-    tensor_model_parallel_all_gather,
-    tensor_model_parallel_reduce_scatter,
-)
-from vllm.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-)
+from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
+                         PassConfig, VllmConfig)
+from vllm.distributed import (tensor_model_parallel_all_gather,
+                              tensor_model_parallel_reduce_scatter)
+from vllm.distributed.parallel_state import (init_distributed_environment,
+                                             initialize_model_parallel)
 from vllm.platforms import current_platform
 from vllm.utils import update_environment_variables
 
 from ..models.registry import HF_EXAMPLE_MODELS
-from ..utils import (
-    compare_two_settings,
-    create_new_process_for_each_test,
-    multi_gpu_test,
-)
+from ..utils import (compare_two_settings, create_new_process_for_each_test,
+                     multi_gpu_test)
 from .backend import TestBackend
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -45,20 +33,21 @@ prompts = [
 
 
 class TestMMRSModel(torch.nn.Module):
+
     def __init__(self, hidden_size=16, dtype=torch.float16):
         super().__init__()
         self.hidden_size = hidden_size
         self.dtype = dtype
-        self.gate_proj = torch.nn.Parameter(
-            torch.empty((self.hidden_size * 2, hidden_size)), requires_grad=False
-        )
+        self.gate_proj = torch.nn.Parameter(torch.empty(
+            (self.hidden_size * 2, hidden_size)),
+                                            requires_grad=False)
         # Initialize weights
         torch.nn.init.normal_(self.gate_proj, std=0.02)
 
     def forward(self, hidden_states):
         """
         Forward pass implementing the mm + reduce scatter in the FX graph
-
+    
         """
         # Reshape input
         view = hidden_states.reshape(-1, self.hidden_size)
@@ -77,13 +66,14 @@ class TestMMRSModel(torch.nn.Module):
 
 
 class TestAGMMModel(torch.nn.Module):
+
     def __init__(self, hidden_size=16, dtype=torch.float16):
         super().__init__()
         self.hidden_size = hidden_size
         self.dtype = dtype
-        self.weight = torch.nn.Parameter(
-            torch.empty((hidden_size, hidden_size)), requires_grad=False
-        )
+        self.weight = torch.nn.Parameter(torch.empty(
+            (hidden_size, hidden_size)),
+                                         requires_grad=False)
         # Initialize weights
         torch.nn.init.normal_(self.weight, std=0.02)
 
@@ -106,35 +96,32 @@ class TestAGMMModel(torch.nn.Module):
 
 
 class _BaseScaledMMModel(torch.nn.Module):
+
     def __init__(self, hidden_size=16, dtype=torch.float16):
         super().__init__()
         self.hidden_size = hidden_size
         self.dtype = dtype
-        self.weight = (
-            torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE)
-            .contiguous()
-            .transpose(0, 1)
-        )
+        self.weight = torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE)\
+            .contiguous().transpose(0, 1)
 
         # Initialize scale_b for _scaled_mm.
         self.scale_b = torch.ones(1, self.hidden_size, dtype=torch.float32)
 
 
 class TestScaledMMRSModel(_BaseScaledMMModel):
+
     def forward(self, input: torch.Tensor):
         """
         Forward pass implementing the scaled_mm + reduce scatter in the FX graph
-
+    
         """
         fp8_input = input.to(FP8_DTYPE)
         scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32)
-        scaled_mm = torch._scaled_mm(
-            fp8_input,
-            self.weight,
-            scale_a=scale_a,
-            scale_b=self.scale_b,
-            out_dtype=self.dtype,
-        )
+        scaled_mm = torch._scaled_mm(fp8_input,
+                                     self.weight,
+                                     scale_a=scale_a,
+                                     scale_b=self.scale_b,
+                                     out_dtype=self.dtype)
         reduce_scatter = tensor_model_parallel_reduce_scatter(scaled_mm, dim=0)
         return reduce_scatter
 
@@ -146,6 +133,7 @@ class TestScaledMMRSModel(_BaseScaledMMModel):
 
 
 class TestAGScaledMMModel(_BaseScaledMMModel):
+
     def forward(self, input: torch.Tensor):
         """
         Forward pass implementing the all gather + scaled_mm in the FX graph
@@ -155,13 +143,11 @@ class TestAGScaledMMModel(_BaseScaledMMModel):
         all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0)
 
         scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32)
-        scaled_mm = torch._scaled_mm(
-            all_gather,
-            self.weight,
-            scale_a=scale_a,
-            scale_b=self.scale_b,
-            out_dtype=self.dtype,
-        )
+        scaled_mm = torch._scaled_mm(all_gather,
+                                     self.weight,
+                                     scale_a=scale_a,
+                                     scale_b=self.scale_b,
+                                     out_dtype=self.dtype)
         return scaled_mm
 
     def ops_in_model_before(self):
@@ -172,22 +158,20 @@ class TestAGScaledMMModel(_BaseScaledMMModel):
 
 
 class TestCutlassScaledMMRSModel(_BaseScaledMMModel):
+
     def forward(self, input: torch.Tensor):
         """
         Forward pass implementing the cutlass_scaled_mm + reduce scatter
         in the FX graph
-
+    
         """
         fp8_input = input.to(FP8_DTYPE)
         scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32)
-        mm_out = torch.empty(
-            (fp8_input.shape[0], self.weight.shape[1]),
-            dtype=self.dtype,
-            device=input.device,
-        )
-        torch.ops._C.cutlass_scaled_mm(
-            mm_out, fp8_input, self.weight, scale_a, self.scale_b, None
-        )
+        mm_out = torch.empty((fp8_input.shape[0], self.weight.shape[1]),
+                             dtype=self.dtype,
+                             device=input.device)
+        torch.ops._C.cutlass_scaled_mm(mm_out, fp8_input, self.weight, scale_a,
+                                       self.scale_b, None)
         reduce_scatter = tensor_model_parallel_reduce_scatter(mm_out, dim=0)
         return reduce_scatter
 
@@ -199,9 +183,10 @@ class TestCutlassScaledMMRSModel(_BaseScaledMMModel):
 
 
 class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
+
     def forward(self, input: torch.Tensor):
         """
-        Forward pass implementing the all gather + cutlass_scaled_mm
+        Forward pass implementing the all gather + cutlass_scaled_mm 
         in the FX graph
         """
         # Reshape input
@@ -210,14 +195,11 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
 
         scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32)
 
-        mm_out = torch.empty(
-            (all_gather.shape[0], self.weight.shape[1]),
-            dtype=self.dtype,
-            device=all_gather.device,
-        )
-        torch.ops._C.cutlass_scaled_mm(
-            mm_out, all_gather, self.weight, scale_a, self.scale_b, None
-        )
+        mm_out = torch.empty((all_gather.shape[0], self.weight.shape[1]),
+                             dtype=self.dtype,
+                             device=all_gather.device)
+        torch.ops._C.cutlass_scaled_mm(mm_out, all_gather, self.weight,
+                                       scale_a, self.scale_b, None)
         return mm_out
 
     def ops_in_model_before(self):
@@ -228,37 +210,23 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
 
 
 @multi_gpu_test(num_gpus=2)
-@pytest.mark.parametrize(
-    "test_model",
-    [
-        TestMMRSModel,
-        TestAGMMModel,
-        TestScaledMMRSModel,
-        TestAGScaledMMModel,
-        TestCutlassScaledMMRSModel,
-        TestAGCutlassScaledMMModel,
-    ],
-)
+@pytest.mark.parametrize("test_model", [
+    TestMMRSModel, TestAGMMModel, TestScaledMMRSModel, TestAGScaledMMModel,
+    TestCutlassScaledMMRSModel, TestAGCutlassScaledMMModel
+])
 @pytest.mark.parametrize("batch_size", [8])
 @pytest.mark.parametrize("seq_len", [16])
 @pytest.mark.parametrize("hidden_size", [16])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
-def test_async_tp_pass_replace(
-    test_model: str, batch_size: int, seq_len: int, hidden_size: int, dtype: torch.dtype
-):
-    if (
-        test_model
-        in (
-            TestScaledMMRSModel,
-            TestAGScaledMMModel,
-            TestCutlassScaledMMRSModel,
-            TestAGCutlassScaledMMModel,
-        )
-        and dtype == torch.float16
-    ):
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
+                    reason="Only test on CUDA")
+def test_async_tp_pass_replace(test_model: str, batch_size: int, seq_len: int,
+                               hidden_size: int, dtype: torch.dtype):
+    if test_model in (TestScaledMMRSModel, TestAGScaledMMModel,
+                      TestCutlassScaledMMRSModel,
+                      TestAGCutlassScaledMMModel) and dtype == torch.float16:
         pytest.skip(
-            "Only bf16 high precision output types are supported for "
+            "Only bf16 high precision output types are supported for " \
             "per-token (row-wise) scaling"
         )
 
@@ -267,24 +235,19 @@ def test_async_tp_pass_replace(
     def run_torch_spawn(fn, nprocs):
         # need to use torch.mp.spawn otherwise will have problems with
         # torch.distributed and cuda
-        torch.multiprocessing.spawn(
-            fn,
-            args=(num_processes, test_model, batch_size, seq_len, hidden_size, dtype),
-            nprocs=nprocs,
-        )
+        torch.multiprocessing.spawn(fn,
+                                    args=(num_processes, test_model,
+                                          batch_size, seq_len, hidden_size,
+                                          dtype),
+                                    nprocs=nprocs)
 
     run_torch_spawn(async_tp_pass_on_test_model, num_processes)
 
 
-def async_tp_pass_on_test_model(
-    local_rank: int,
-    world_size: int,
-    test_model_cls: torch.nn.Module,
-    batch_size: int,
-    seq_len: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-):
+def async_tp_pass_on_test_model(local_rank: int, world_size: int,
+                                test_model_cls: torch.nn.Module,
+                                batch_size: int, seq_len: int,
+                                hidden_size: int, dtype: torch.dtype):
     current_platform.seed_everything(0)
 
     device = torch.device(f"cuda:{local_rank}")
@@ -292,15 +255,13 @@ def async_tp_pass_on_test_model(
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "12345",
-        }
-    )
+    update_environment_variables({
+        'RANK': str(local_rank),
+        'LOCAL_RANK': str(local_rank),
+        'WORLD_SIZE': str(world_size),
+        'MASTER_ADDR': 'localhost',
+        'MASTER_PORT': '12345',
+    })
 
     # initialize distributed
     init_distributed_environment()
@@ -308,28 +269,27 @@ def async_tp_pass_on_test_model(
 
     # configure vllm config for SequenceParallelismPass
     vllm_config = VllmConfig()
-    vllm_config.compilation_config = CompilationConfig(
-        pass_config=PassConfig(
-            enable_async_tp=True,
-        ),
-    )
+    vllm_config.compilation_config = CompilationConfig(pass_config=PassConfig(
+        enable_async_tp=True, ), )
     vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.
     model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
-    vllm_config.model_config = ModelConfig(
-        model=model_name, trust_remote_code=True, dtype=dtype, seed=42
-    )
+    vllm_config.model_config = ModelConfig(model=model_name,
+                                           trust_remote_code=True,
+                                           dtype=dtype,
+                                           seed=42)
 
     async_tp_pass = AsyncTPPass(vllm_config)
     backend = TestBackend(async_tp_pass)
 
-    model = test_model_cls(hidden_size, dtype)  # Pass dtype to model constructor
+    model = test_model_cls(hidden_size,
+                           dtype)  # Pass dtype to model constructor
 
-    hidden_states = torch.randn(
-        (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False
-    )
+    hidden_states = torch.randn((batch_size * seq_len, hidden_size),
+                                dtype=dtype,
+                                requires_grad=False)
 
     compiled_model = torch.compile(model, backend=backend)
     compiled_model(hidden_states)
@@ -346,10 +306,10 @@ def async_tp_pass_on_test_model(
 
 
 @create_new_process_for_each_test()
-@pytest.mark.parametrize(
-    "model_id",
-    ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"],
-)
+@pytest.mark.parametrize("model_id", [
+    "meta-llama/Llama-3.2-1B-Instruct",
+    "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"
+])
 @pytest.mark.parametrize("tp_size", [2])
 @pytest.mark.parametrize("async_tp_enabled", [True])
 @pytest.mark.parametrize("distributed_backend", ["mp"])
@@ -382,10 +342,12 @@ def test_async_tp_pass_correctness(
         common_args.append("--enforce-eager")
 
     compilation_config = {
-        "level": 3,
-        "compile_sizes": [2, 4, 8],
-        "splitting_ops": [],
-        "pass_config": {"enable_async_tp": async_tp_enabled},
+        'level': 3,
+        'compile_sizes': [2, 4, 8],
+        'splitting_ops': [],
+        'pass_config': {
+            'enable_async_tp': async_tp_enabled
+        },
     }
 
     async_tp_env = tp_env = {
@@ -410,6 +372,9 @@ def test_async_tp_pass_correctness(
         "mp",
     ]
 
-    compare_two_settings(
-        model_id, async_tp_args, tp_args, async_tp_env, tp_env, method="generate"
-    )
+    compare_two_settings(model_id,
+                         async_tp_args,
+                         tp_args,
+                         async_tp_env,
+                         tp_env,
+                         method="generate")

tests/compile/test_basic_correctness.py

@@ -103,28 +103,23 @@ def test_compile_correctness(
     attn_backend = test_setting.attn_backend
     method = test_setting.method
     if cuda_device_count_stateless() < pp_size * tp_size:
-        pytest.skip(
-            f"Need at least {pp_size}*{tp_size} CUDA gpus but got "
-            f"{cuda_device_count_stateless()}"
-        )
+        pytest.skip(f"Need at least {pp_size}*{tp_size} CUDA gpus but got "
+                    f"{cuda_device_count_stateless()}")
 
     with monkeypatch.context() as m:
         m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
         final_args = [
-            "--enforce-eager",
-            *model_args,
-            "-pp",
-            str(pp_size),
-            "-tp",
-            str(tp_size),
+            "--enforce-eager", *model_args, "-pp",
+            str(pp_size), "-tp",
+            str(tp_size)
         ]
 
         all_args: list[list[str]] = []
         all_envs: list[dict[str, str] | None] = []
 
         for level in [
-            CompilationLevel.NO_COMPILATION,
-            CompilationLevel.PIECEWISE,
+                CompilationLevel.NO_COMPILATION,
+                CompilationLevel.PIECEWISE,
         ]:
             all_args.append(final_args + [f"-O{level}"])
             all_envs.append({})
@@ -135,15 +130,14 @@ def test_compile_correctness(
             model,
             all_args,
             all_envs,
-            method=method if method != "generate" else "generate_close",
-        )
+            method=method if method != "generate" else "generate_close")
         all_envs.clear()
         all_args.clear()
 
         for level in [
-            CompilationLevel.NO_COMPILATION,
-            CompilationLevel.DYNAMO_AS_IS,
-            CompilationLevel.DYNAMO_ONCE,
+                CompilationLevel.NO_COMPILATION,
+                CompilationLevel.DYNAMO_AS_IS,
+                CompilationLevel.DYNAMO_ONCE,
         ]:
             all_args.append(final_args + [f"-O{level}"])
             all_envs.append({})

tests/compile/test_config.py

@@ -9,11 +9,11 @@ from vllm.utils import _is_torch_equal_or_newer
 
 
 def test_version():
-    assert _is_torch_equal_or_newer("2.8.0.dev20250624+cu128", "2.8.0.dev")
-    assert _is_torch_equal_or_newer("2.8.0a0+gitc82a174", "2.8.0.dev")
-    assert _is_torch_equal_or_newer("2.8.0", "2.8.0.dev")
-    assert _is_torch_equal_or_newer("2.8.1", "2.8.0.dev")
-    assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev")
+    assert _is_torch_equal_or_newer('2.8.0.dev20250624+cu128', '2.8.0.dev')
+    assert _is_torch_equal_or_newer('2.8.0a0+gitc82a174', '2.8.0.dev')
+    assert _is_torch_equal_or_newer('2.8.0', '2.8.0.dev')
+    assert _is_torch_equal_or_newer('2.8.1', '2.8.0.dev')
+    assert not _is_torch_equal_or_newer('2.7.1', '2.8.0.dev')
 
 
 def test_use_cudagraphs_dynamic(monkeypatch):
@@ -21,7 +21,7 @@ def test_use_cudagraphs_dynamic(monkeypatch):
     vllm_config = VllmConfig()
     assert vllm_config.compilation_config.use_cudagraph
 
-    monkeypatch.setenv("VLLM_USE_V1", "0")
+    monkeypatch.setenv('VLLM_USE_V1', '0')
     vllm_config = VllmConfig()
     assert not vllm_config.compilation_config.use_cudagraph
 
@@ -44,23 +44,19 @@ def test_VLLM_DISABLE_COMPILE_CACHE(vllm_runner, monkeypatch, val):
     assert vllm.envs.VLLM_USE_V1
 
     # Disable multiprocessing so that the counter is in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
-    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", val)
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
+    monkeypatch.setenv('VLLM_DISABLE_COMPILE_CACHE', val)
 
     compilation_config = {
         "use_cudagraph": False,  # speed things up a bit
     }
     with (
-        compilation_counter.expect(
-            num_cache_entries_updated=0, num_compiled_artifacts_saved=0
-        ),
-        # loading the model causes compilation (if enabled) to happen
-        vllm_runner(
-            "facebook/opt-125m",
-            compilation_config=compilation_config,
-            gpu_memory_utilization=0.4,
-        ) as _,
-    ):
+            compilation_counter.expect(num_cache_entries_updated=0,
+                                       num_compiled_artifacts_saved=0),
+            # loading the model causes compilation (if enabled) to happen
+            vllm_runner('facebook/opt-125m',
+                        compilation_config=compilation_config,
+                        gpu_memory_utilization=0.4) as _):
         pass
 
 
@@ -71,25 +67,22 @@ def test_use_cudagraphs(vllm_runner, monkeypatch, enabled):
     assert vllm.envs.VLLM_USE_V1
 
     # Disable multiprocessing so that the counter is in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
 
     compilation_config = {
         "cudagraph_capture_sizes": [100],
         "use_cudagraph": enabled,
     }
     with (
-        compilation_counter.expect(
-            num_graphs_seen=1,
-            num_gpu_runner_capture_triggers=1 if enabled else 0,
-            num_cudagraph_captured=13 if enabled else 0,
-        ),
-        # loading the model causes compilation (if enabled) to happen
-        vllm_runner(
-            "facebook/opt-125m",
-            compilation_config=compilation_config,
-            gpu_memory_utilization=0.4,
-        ) as _,
-    ):
+            compilation_counter.expect(
+                num_graphs_seen=1,
+                num_gpu_runner_capture_triggers=1 if enabled else 0,
+                num_cudagraph_captured=13 if enabled else 0,
+            ),
+            # loading the model causes compilation (if enabled) to happen
+            vllm_runner('facebook/opt-125m',
+                        compilation_config=compilation_config,
+                        gpu_memory_utilization=0.4) as _):
         pass
 
 
@@ -97,17 +90,14 @@ def test_use_cudagraphs(vllm_runner, monkeypatch, enabled):
 @pytest.mark.forked
 def test_dynamo_as_is(vllm_runner, monkeypatch):
     # Disable multiprocessing so that the counter is in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
 
     with (
-        compilation_counter.expect(dynamo_as_is_count=1),
-        # loading the model causes compilation (if enabled) to happen
-        vllm_runner(
-            "facebook/opt-125m",
-            compilation_config={"level": 1},
-            gpu_memory_utilization=0.4,
-        ) as _,
-    ):
+            compilation_counter.expect(dynamo_as_is_count=1),
+            # loading the model causes compilation (if enabled) to happen
+            vllm_runner('facebook/opt-125m',
+                        compilation_config={"level": 1},
+                        gpu_memory_utilization=0.4) as _):
         pass
 
 
@@ -115,16 +105,14 @@ def test_dynamo_as_is(vllm_runner, monkeypatch):
 @pytest.mark.forked
 def test_no_compilation(vllm_runner, monkeypatch):
     # Disable multiprocessing so that the counter is in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
     with (
-        compilation_counter.expect(num_graphs_seen=0, dynamo_as_is_count=0),
-        # loading the model causes compilation (if enabled) to happen
-        vllm_runner(
-            "facebook/opt-125m",
-            compilation_config={"level": 0},
-            gpu_memory_utilization=0.4,
-        ) as _,
-    ):
+            compilation_counter.expect(num_graphs_seen=0,
+                                       dynamo_as_is_count=0),
+            # loading the model causes compilation (if enabled) to happen
+            vllm_runner('facebook/opt-125m',
+                        compilation_config={"level": 0},
+                        gpu_memory_utilization=0.4) as _):
         pass
 
 
@@ -132,73 +120,77 @@ def test_no_compilation(vllm_runner, monkeypatch):
 @pytest.mark.forked
 def test_enforce_eager(vllm_runner, monkeypatch):
     # Disable multiprocessing so that the counter is in the same process
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', '0')
 
     with (
-        compilation_counter.expect(num_graphs_seen=0, dynamo_as_is_count=0),
-        # loading the model causes compilation (if enabled) to happen
-        vllm_runner(
-            "facebook/opt-125m", enforce_eager=True, gpu_memory_utilization=0.4
-        ) as _,
-    ):
+            compilation_counter.expect(num_graphs_seen=0,
+                                       dynamo_as_is_count=0),
+            # loading the model causes compilation (if enabled) to happen
+            vllm_runner('facebook/opt-125m',
+                        enforce_eager=True,
+                        gpu_memory_utilization=0.4) as _):
         pass
 
 
 def test_splitting_ops_dynamic():
     # Default config
     config = VllmConfig()
-    assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
+    assert config.compilation_config.cudagraph_mode == \
+        CUDAGraphMode.FULL_AND_PIECEWISE
     assert config.compilation_config.splitting_ops_contain_attention()
 
     # When use_inductor_graph_partition=True
-    if _is_torch_equal_or_newer("2.9.0.dev"):
+    if _is_torch_equal_or_newer('2.9.0.dev'):
         # inductor graph partition is only available in PyTorch 2.9+.
         # this is a fast config check so we are not using pytest.skip.
-        config = VllmConfig(
-            compilation_config=CompilationConfig(
-                use_inductor_graph_partition=True, splitting_ops=["silly_attention"]
-            )
-        )
+        config = VllmConfig(compilation_config=CompilationConfig(
+            use_inductor_graph_partition=True,
+            splitting_ops=["silly_attention"]))
         # should ignore splitting_ops
         assert config.compilation_config.splitting_ops == []
 
     # When attn_fusion pass enabled.
-    config = VllmConfig(
-        compilation_config=CompilationConfig(
-            pass_config={"enable_attn_fusion": True, "enable_noop": True},
-            custom_ops=["+quant_fp8"],
-            cudagraph_mode=CUDAGraphMode.PIECEWISE,
-        )
-    )
+    config = VllmConfig(compilation_config=CompilationConfig(
+        pass_config={
+            "enable_attn_fusion": True,
+            "enable_noop": True
+        },
+        custom_ops=["+quant_fp8"],
+        cudagraph_mode=CUDAGraphMode.PIECEWISE,
+    ))
     assert config.compilation_config.splitting_ops == []
     # cudagraph mode also fall back to FULL
-    assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL
+    assert config.compilation_config.cudagraph_mode == \
+        CUDAGraphMode.FULL
 
     # splitting_ops can not contain attention ops when attn_fusion
     # pass enabled.
     with pytest.raises(AssertionError):
-        config = VllmConfig(
-            compilation_config=CompilationConfig(
-                pass_config={"enable_attn_fusion": True, "enable_noop": True},
-                custom_ops=["+quant_fp8"],
-                cudagraph_mode=CUDAGraphMode.PIECEWISE,
-                # work around for accessing all attntion ops
-                splitting_ops=CompilationConfig()._attention_ops,
-            )
-        )
+        config = VllmConfig(compilation_config=CompilationConfig(
+            pass_config={
+                "enable_attn_fusion": True,
+                "enable_noop": True
+            },
+            custom_ops=["+quant_fp8"],
+            cudagraph_mode=CUDAGraphMode.PIECEWISE,
+            # work around for accessing all attntion ops
+            splitting_ops=CompilationConfig()._attention_ops,
+        ))
 
     # When both use_inductor_graph_partition and attn_fusion pass enabled.
-    if _is_torch_equal_or_newer("2.9.0.dev"):
-        config = VllmConfig(
-            compilation_config=CompilationConfig(
-                use_inductor_graph_partition=True,
-                pass_config={"enable_attn_fusion": True, "enable_noop": True},
-                custom_ops=["+quant_fp8"],
-                cudagraph_mode=CUDAGraphMode.PIECEWISE,
-            )
-        )
+    if _is_torch_equal_or_newer('2.9.0.dev'):
+        config = VllmConfig(compilation_config=CompilationConfig(
+            use_inductor_graph_partition=True,
+            pass_config={
+                "enable_attn_fusion": True,
+                "enable_noop": True
+            },
+            custom_ops=["+quant_fp8"],
+            cudagraph_mode=CUDAGraphMode.PIECEWISE,
+        ))
         assert config.compilation_config.splitting_ops == []
         # enable_attn_fusion is directly support under
         # use_inductor_graph_partition=True, and cudagraph_mode
         # is unchanged.
-        assert config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE
+        assert config.compilation_config.cudagraph_mode == \
+            CUDAGraphMode.PIECEWISE

tests/compile/test_decorator.py

@@ -4,15 +4,10 @@ import torch
 from torch import nn
 
 from vllm.compilation.counter import compilation_counter
-from vllm.compilation.decorators import ignore_torch_compile, support_torch_compile
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    CompilationLevel,
-    CUDAGraphMode,
-    VllmConfig,
-    set_current_vllm_config,
-)
+from vllm.compilation.decorators import (ignore_torch_compile,
+                                         support_torch_compile)
+from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
+                         CUDAGraphMode, VllmConfig, set_current_vllm_config)
 from vllm.forward_context import BatchDescriptor, set_forward_context
 
 # This import automatically registers `torch.ops.silly.attention`
@@ -23,42 +18,32 @@ MLP_SIZE = 128
 
 
 @torch.inference_mode
-def run_model(
-    vllm_config: VllmConfig, model: nn.Module, cudagraph_runtime_mode: CUDAGraphMode
-):
+def run_model(vllm_config: VllmConfig, model: nn.Module,
+              cudagraph_runtime_mode: CUDAGraphMode):
     with set_forward_context({}, vllm_config=vllm_config):
         # warmup for the model with cudagraph_mode NONE
         model(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
 
         # simulate cudagraphs capturing
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             model(torch.randn(2, MLP_SIZE).cuda())
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=1,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=1, )):
             model(torch.randn(1, MLP_SIZE).cuda())
 
         # simulate cudagraphs replay
-        with set_forward_context(
-            {},
-            vllm_config=vllm_config,
-            cudagraph_runtime_mode=cudagraph_runtime_mode,
-            batch_descriptor=BatchDescriptor(
-                num_tokens=2,
-            ),
-        ):
+        with set_forward_context({},
+                                 vllm_config=vllm_config,
+                                 cudagraph_runtime_mode=cudagraph_runtime_mode,
+                                 batch_descriptor=BatchDescriptor(
+                                     num_tokens=2, )):
             output = model(torch.randn(2, MLP_SIZE).cuda())
 
         output = output.cpu()
@@ -67,21 +52,22 @@ def run_model(
 
 def test_ignore_torch_compile_decorator():
     # piecewise
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=True,
-            splitting_ops=["silly.attention"],
-            cudagraph_capture_sizes=[1, 2],
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=True,
+        splitting_ops=["silly.attention"],
+        cudagraph_capture_sizes=[1, 2],
+    ))
     cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
 
     @support_torch_compile
     class A(nn.Module):
-        def __init__(
-            self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs
-        ) -> None:
+
+        def __init__(self,
+                     *,
+                     vllm_config: VllmConfig,
+                     prefix: str = '',
+                     **kwargs) -> None:
             super().__init__()
 
         def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -93,60 +79,66 @@ def test_ignore_torch_compile_decorator():
             return x
 
     @ignore_torch_compile
-    class B(A): ...
+    class B(A):
+        ...
 
     @support_torch_compile
-    class C(B): ...
+    class C(B):
+        ...
 
     with set_current_vllm_config(vllm_config):
-        mod_A = A(vllm_config=vllm_config, prefix="").eval().cuda()
+        mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
 
     # A has support_torch_compile
     with compilation_counter.expect(
-        num_graphs_seen=1,
-        num_piecewise_graphs_seen=3,
-        num_piecewise_capturable_graphs_seen=2,
-        num_backend_compilations=2,
-        num_cudagraph_captured=4,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=1,
+            num_piecewise_graphs_seen=3,
+            num_piecewise_capturable_graphs_seen=2,
+            num_backend_compilations=2,
+            num_cudagraph_captured=4,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
         run_model(vllm_config, mod_A, cudagraph_runtime_mode)
 
     with set_current_vllm_config(vllm_config):
-        mod_B = B(vllm_config=vllm_config, prefix="").eval().cuda()
+        mod_B = B(vllm_config=vllm_config, prefix='').eval().cuda()
 
     # B's ignore_torch_compile should override A's support_torch_compile
     with compilation_counter.expect(
-        num_graphs_seen=0,
-        num_piecewise_graphs_seen=0,
-        num_piecewise_capturable_graphs_seen=0,
-        num_backend_compilations=0,
-        num_cudagraph_captured=0,
+            num_graphs_seen=0,
+            num_piecewise_graphs_seen=0,
+            num_piecewise_capturable_graphs_seen=0,
+            num_backend_compilations=0,
+            num_cudagraph_captured=0,
     ):
         run_model(vllm_config, mod_B, cudagraph_runtime_mode)
 
     with set_current_vllm_config(vllm_config):
-        mod_C = C(vllm_config=vllm_config, prefix="").eval().cuda()
+        mod_C = C(vllm_config=vllm_config, prefix='').eval().cuda()
 
     # C's support_torch_compile should override B's ignore_torch_compile
     with compilation_counter.expect(
-        num_graphs_seen=1,
-        num_piecewise_graphs_seen=3,
-        num_piecewise_capturable_graphs_seen=2,
-        num_backend_compilations=2,
-        num_cudagraph_captured=4,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=1,
+            num_piecewise_graphs_seen=3,
+            num_piecewise_capturable_graphs_seen=2,
+            num_backend_compilations=2,
+            num_cudagraph_captured=4,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
         run_model(vllm_config, mod_C, cudagraph_runtime_mode)
 
 
 # Only enable torch.compile if
 # vllm_config.cache_config.kv_sharing_fast_prefill=True
-@support_torch_compile(
-    enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill
-)
+@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.
+                       kv_sharing_fast_prefill)
 class B(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -160,11 +152,15 @@ class B(nn.Module):
 
 # Only enable torch.compile if
 # vllm_config.cache_config.kv_sharing_fast_prefill=False
-@support_torch_compile(
-    enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
-)
+@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.
+                       cache_config.kv_sharing_fast_prefill)
 class A(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
         super().__init__()
         self.mod1 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
         self.mod2 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
@@ -179,60 +175,54 @@ class A(nn.Module):
 
 
 def test_conditional_compile_enable_if():
-    vllm_config = VllmConfig(
-        cache_config=CacheConfig(
-            kv_sharing_fast_prefill=True,
-        ),
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=True,
-            splitting_ops=["silly.attention"],
-            cudagraph_capture_sizes=[1, 2],
-        ),
-    )
+    vllm_config = VllmConfig(cache_config=CacheConfig(
+        kv_sharing_fast_prefill=True, ),
+                             compilation_config=CompilationConfig(
+                                 level=CompilationLevel.PIECEWISE,
+                                 use_cudagraph=True,
+                                 splitting_ops=["silly.attention"],
+                                 cudagraph_capture_sizes=[1, 2],
+                             ))
     cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
 
     with set_current_vllm_config(vllm_config):
-        mod_A = A(vllm_config=vllm_config, prefix="").eval().cuda()
+        mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
 
     # A has support_torch_compile but enable_if fn returns False
     # enalbe_if will be True for B, so we expect mod1 and mod2
     # to be compiled
     with compilation_counter.expect(
-        num_graphs_seen=2,
-        num_piecewise_graphs_seen=6,
-        # 3 piecewise graphs per instance of B()
-        num_piecewise_capturable_graphs_seen=4,
-        num_backend_compilations=4,
-        num_cudagraph_captured=8,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=2,
+            num_piecewise_graphs_seen=6,
+            # 3 piecewise graphs per instance of B()
+            num_piecewise_capturable_graphs_seen=4,
+            num_backend_compilations=4,
+            num_cudagraph_captured=8,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
         run_model(vllm_config, mod_A, cudagraph_runtime_mode)
 
     # Set kv_sharing_fast_prefill=False
     # which will cause A to be compiled and B to not be compiled
-    vllm_config = VllmConfig(
-        cache_config=CacheConfig(
-            kv_sharing_fast_prefill=False,
-        ),
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            use_cudagraph=True,
-            splitting_ops=["silly.attention"],
-            cudagraph_capture_sizes=[1, 2],
-        ),
-    )
+    vllm_config = VllmConfig(cache_config=CacheConfig(
+        kv_sharing_fast_prefill=False, ),
+                             compilation_config=CompilationConfig(
+                                 level=CompilationLevel.PIECEWISE,
+                                 use_cudagraph=True,
+                                 splitting_ops=["silly.attention"],
+                                 cudagraph_capture_sizes=[1, 2],
+                             ))
 
     with set_current_vllm_config(vllm_config):
-        mod_A = A(vllm_config=vllm_config, prefix="").eval().cuda()
+        mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
 
     with compilation_counter.expect(
-        num_graphs_seen=1,
-        num_piecewise_graphs_seen=7,
-        # 3 attn ops and 4 non-attn ops
-        num_piecewise_capturable_graphs_seen=4,
-        num_backend_compilations=4,
-        num_cudagraph_captured=8,
-        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+            num_graphs_seen=1,
+            num_piecewise_graphs_seen=7,
+            # 3 attn ops and 4 non-attn ops
+            num_piecewise_capturable_graphs_seen=4,
+            num_backend_compilations=4,
+            num_cudagraph_captured=8,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
     ):
         run_model(vllm_config, mod_A, cudagraph_runtime_mode)

tests/compile/test_full_graph.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import logging
 import tempfile
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import pytest
 import torch
@@ -14,64 +14,54 @@ from tests.quantization.utils import is_quant_method_supported
 from vllm import LLM, SamplingParams
 from vllm.attention.backends.registry import _Backend
 from vllm.attention.selector import global_force_attn_backend_context_manager
-from vllm.config import CompilationConfig, CompilationLevel, CUDAGraphMode, PassConfig
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
+                         PassConfig)
 from vllm.platforms import current_platform
 from vllm.utils import is_torch_equal_or_newer
 
 from ..utils import create_new_process_for_each_test
 
 
-def models_list(*, all: bool = True, keywords: list[str] | None = None):
+def models_list(*, all: bool = True, keywords: Optional[list[str]] = None):
     TEST_MODELS: list[tuple[str, dict[str, Any]]] = [
         ("facebook/opt-125m", {}),
-        (
-            "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
-            {
-                "dtype": torch.float16,
-            },
-        ),
-        (
-            "neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic",
-            {
-                "dtype": torch.float16,
-            },
-        ),
+        ("nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", {
+            "dtype": torch.float16,
+        }),
+        ("neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic", {
+            "dtype": torch.float16,
+        }),
         ("neuralmagic/Llama-3.2-1B-Instruct-quantized.w8a8", {}),
         ("meta-llama/Llama-3.2-1B-Instruct", {}),
     ]
 
     if all:
+
         # TODO: figure out why this fails.
         if False and is_quant_method_supported("gguf"):  # noqa: SIM223
-            TEST_MODELS.append(
-                ("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", {"quantization": "gguf"})
-            )
+            TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", {
+                "quantization": "gguf"
+            }))
 
         if is_quant_method_supported("gptq"):
-            TEST_MODELS.append(
-                ("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {"quantization": "gptq"})
-            )
+            TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", {
+                "quantization": "gptq"
+            }))
 
         if is_quant_method_supported("gptq_marlin"):
-            TEST_MODELS.append(
-                (
-                    "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",
-                    {"quantization": "gptq_marlin"},
-                )
-            )
+            TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", {
+                "quantization": "gptq_marlin"
+            }))
 
         if is_quant_method_supported("gptq_marlin_24"):
-            TEST_MODELS.append(
-                (
-                    "alexm-nm/tinyllama-24-marlin24-4bit-g128",
-                    {"quantization": "gptq_marlin_24"},
-                )
-            )
+            TEST_MODELS.append(("alexm-nm/tinyllama-24-marlin24-4bit-g128", {
+                "quantization": "gptq_marlin_24"
+            }))
 
         if not current_platform.is_rocm() and is_quant_method_supported("awq"):
-            TEST_MODELS.append(
-                ("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {"quantization": "AWQ"})
-            )
+            TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {
+                "quantization": "AWQ"
+            }))
 
     if keywords is None:
         return TEST_MODELS
@@ -105,34 +95,22 @@ def test_full_graph(
     "compilation_config, model_info",
     [
         # additional compile sizes, only some of the models
-        (
-            CompilationConfig(level=CompilationLevel.PIECEWISE, compile_sizes=[1, 2]),
-            model,
-        )
+        (CompilationConfig(level=CompilationLevel.PIECEWISE,
+                           compile_sizes=[1, 2]), model)
         for model in models_list(all=False)
-    ]
-    + [
+    ] + [
         # RMSNorm + quant fusion, only 8-bit quant models
-        (
-            CompilationConfig(
-                level=CompilationLevel.PIECEWISE,
-                custom_ops=["+rms_norm"],
-                pass_config=PassConfig(enable_fusion=True, enable_noop=True),
-            ),
-            model,
-        )
+        (CompilationConfig(level=CompilationLevel.PIECEWISE,
+                           custom_ops=["+rms_norm"],
+                           pass_config=PassConfig(enable_fusion=True,
+                                                  enable_noop=True)), model)
         for model in models_list(keywords=["FP8-dynamic", "quantized.w8a8"])
-    ]
-    + [
+    ] + [
         # Test depyf integration works
-        (
-            CompilationConfig(
-                level=CompilationLevel.PIECEWISE, debug_dump_path=tempfile.gettempdir()
-            ),
-            ("facebook/opt-125m", {}),
-        ),
-    ]
-    + [
+        (CompilationConfig(level=CompilationLevel.PIECEWISE,
+                           debug_dump_path=tempfile.gettempdir()),
+         ("facebook/opt-125m", {})),
+    ] + [
         # graph inductor partition
         (
             CompilationConfig(
@@ -141,24 +119,20 @@ def test_full_graph(
                 # torch._C.Tag.cudagraph_unsafe to specify splitting ops
                 use_inductor_graph_partition=True,
                 cudagraph_mode=CUDAGraphMode.PIECEWISE,
-                compile_sizes=[1, 2],
-            ),
-            model,
-        )
-        for model in models_list(all=False)
+                compile_sizes=[1, 2]),
+            model) for model in models_list(all=False)
         if is_torch_equal_or_newer("2.9.0.dev")
-    ],
-)
+    ])
 # only test some of the models
 @create_new_process_for_each_test()
 def test_custom_compile_config(
     compilation_config: CompilationConfig,
     model_info: tuple[str, dict[str, Any]],
 ):
-    if compilation_config.use_inductor_graph_partition and not is_torch_equal_or_newer(
-        "2.9.0.dev"
-    ):
-        pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
+    if (compilation_config.use_inductor_graph_partition
+            and not is_torch_equal_or_newer("2.9.0.dev")):
+        pytest.skip("inductor graph partition is only available "
+                    "in PyTorch 2.9+")
 
     model, model_kwargs = model_info
     print(f"MODEL={model}")
@@ -182,7 +156,8 @@ def test_fp8_kv_scale_compile(optimization_level: int):
 
 def test_inductor_graph_partition_attn_fusion(caplog_vllm):
     if not is_torch_equal_or_newer("2.9.0.dev"):
-        pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
+        pytest.skip("inductor graph partition is only available "
+                    "in PyTorch 2.9+")
 
     model = "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
     compilation_config = CompilationConfig(
@@ -196,16 +171,14 @@ def test_inductor_graph_partition_attn_fusion(caplog_vllm):
         "kv_cache_dtype": "fp8",
         "max_model_len": 1024,
     }
-    with (
-        caplog_vllm.at_level(logging.DEBUG),
-        global_force_attn_backend_context_manager(_Backend.FLASHINFER),
-    ):
+    with caplog_vllm.at_level(
+            logging.DEBUG), global_force_attn_backend_context_manager(
+                _Backend.FLASHINFER):
         run_model(compilation_config, model, model_kwargs)
 
     try:
-        assert "Fused quantization onto 48 attention nodes" in caplog_vllm.text, (
-            caplog_vllm.text
-        )
+        assert ("Fused quantization onto 48 attention nodes"
+                in caplog_vllm.text), caplog_vllm.text
     except AssertionError:
         # Note: this message is only triggered when the compilation goes
         # through the custom pass. Due to multiple layers of cache on
@@ -216,11 +189,8 @@ def test_inductor_graph_partition_attn_fusion(caplog_vllm):
         assert "Fused quantization" not in caplog_vllm.text
 
 
-def run_model(
-    compile_config: Union[int, CompilationConfig],
-    model: str,
-    model_kwargs: dict[str, Any],
-):
+def run_model(compile_config: Union[int, CompilationConfig], model: str,
+              model_kwargs: dict[str, Any]):
     prompts = [
         "Hello, my name is",
         "The president of the United States is",

tests/compile/test_functionalization.py

@@ -5,254 +5,114 @@ import pytest
 import torch
 
 import vllm.envs as envs
+from vllm import LLM, SamplingParams
 from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
 from vllm.compilation.fix_functionalization import FixFunctionalizationPass
-from vllm.compilation.fusion import RMSNormQuantFusionPass
+from vllm.compilation.fusion import FUSED_OPS, RMSNormQuantFusionPass
 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
 from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
 from vllm.config import CompilationConfig, PassConfig, VllmConfig
-from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
-from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.platforms import current_platform
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey, kFp8DynamicTokenSym, kFp8StaticTensorSym)
 
 from .backend import TestBackend
 
-TEST_FP8 = current_platform.supports_fp8()
-FP8_DTYPE = current_platform.fp8_dtype()
+OPS_IN_MODEL = [
+    torch.ops._C.rotary_embedding.default,
+    torch.ops._C.fused_add_rms_norm.default,
+]
 
+RMS_OP = torch.ops._C.rms_norm.default
 
-class TestSiluMul(torch.nn.Module):
-    def __init__(self, hidden_size: int = 128):
-        super().__init__()
-        self.silu_and_mul = SiluAndMul()
-        self.wscale = torch.rand(1, dtype=torch.float32)
-        self.scale = torch.rand(1, dtype=torch.float32)
+RMS_QUANT_OPS = {
+    "static_fp8": [
+        torch.ops._C.rms_norm_static_fp8_quant.default,
+        torch.ops._C.fused_add_rms_norm_static_fp8_quant.default
+    ],
+}
 
-        if TEST_FP8:
-            self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
-            self.fp8_linear = Fp8LinearOp(
-                act_quant_static=True,
-                act_quant_group_shape=GroupShape.PER_TENSOR,
-            )
+SILU_MUL_OP = torch.ops._C.silu_and_mul.default
 
-    def forward(self, x):
-        y = self.silu_and_mul(x)
-        if TEST_FP8:
-            x2 = self.fp8_linear.apply(y, self.w, self.wscale, input_scale=self.wscale)
-            return x2
-        else:
-            return y
-
-    def example_inputs(self, num_tokens=32, hidden_size=128):
-        dtype = torch.float16 if TEST_FP8 else torch.float32
-        return (torch.rand(num_tokens, hidden_size * 2, dtype=dtype),)
-
-    def ops_in_model(self, do_fusion):
-        if TEST_FP8 and do_fusion:
-            return [torch.ops._C.silu_and_mul_quant.default]
-        else:
-            return [torch.ops._C.silu_and_mul.default]
-
-    def ops_not_in_model(self):
-        return []
-
-
-class TestFusedAddRMSNorm(torch.nn.Module):
-    def __init__(self, hidden_size=16, intermediate_size=32):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-
-        dtype = torch.float16 if TEST_FP8 else torch.float32
-
-        self.gate_proj = torch.nn.Parameter(
-            torch.empty((intermediate_size, hidden_size), dtype=dtype)
-        )
-        self.norm = RMSNorm(intermediate_size, 1e-05)
-        self.norm.weight = torch.nn.Parameter(
-            torch.ones(intermediate_size, dtype=dtype)
-        )
-
-        torch.nn.init.normal_(self.gate_proj, std=0.02)
-
-        if TEST_FP8:
-            self.fp8_linear = Fp8LinearOp(act_quant_static=True)
-
-            self.scale = torch.rand(1, dtype=torch.float32)
-            self.w = torch.rand(hidden_size, intermediate_size).to(dtype=FP8_DTYPE).t()
-            self.wscale = torch.rand(1, dtype=torch.float32)
-
-    def forward(self, hidden_states, residual):
-        # Reshape input
-        view = hidden_states.reshape(-1, self.hidden_size)
-
-        # matrix multiplication
-        permute = self.gate_proj.permute(1, 0)
-        mm = torch.mm(view, permute)
-
-        # layer normalization
-        norm_output, residual_output = self.norm(mm, residual)
-
-        if TEST_FP8:
-            # scaled_mm with static input quantization
-            fp8_linear_result = self.fp8_linear.apply(
-                norm_output,
-                self.w,
-                self.wscale,
-                input_scale=self.scale.to(norm_output.device),
-            )
-
-            return fp8_linear_result, residual_output
-
-        else:
-            return norm_output, residual_output
-
-    def example_inputs(self, batch_size=8, hidden_size=16, seq_len=16):
-        dtype = torch.float16 if TEST_FP8 else torch.float32
-        hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
-        residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
-        return (hidden_states, residual)
-
-    def ops_in_model(self, do_fusion):
-        if TEST_FP8 and do_fusion:
-            return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
-        else:
-            return [torch.ops._C.fused_add_rms_norm.default]
-
-    def ops_not_in_model(self):
-        return []
-
-
-class TestRotaryEmbedding(torch.nn.Module):
-    def __init__(self, head_dim=64, rotary_dim=None, max_position=2048, base=10000):
-        super().__init__()
-        self.head_dim = head_dim
-        self.rotary_dim = rotary_dim or head_dim
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            rotary_dim=self.rotary_dim,
-            max_position=max_position,
-            base=base,
-        )
-
-    def forward(self, positions, q, k):
-        q_rotated, k_rotated = self.rotary_emb(positions, q, k)
-        return q_rotated, k_rotated
-
-    def example_inputs(self, num_tokens=32, head_dim=64):
-        dtype = torch.float16
-        positions = torch.arange(num_tokens, dtype=torch.long)
-        q = torch.randn(num_tokens, head_dim, dtype=dtype)
-        k = torch.randn(num_tokens, head_dim, dtype=dtype)
-        return (positions, q, k)
-
-    def ops_in_model(self, do_fusion):
-        return [torch.ops._C.rotary_embedding.default]
-
-    def ops_not_in_model(self):
-        return []
-
-
-class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
-    def __init__(self, head_dim=64, num_heads=4, max_position=2048, base=10000):
-        super().__init__()
-        self.head_dim = head_dim
-        self.num_heads = num_heads
-        self.hidden_size = head_dim * num_heads
-
-        self.qkv_proj = torch.nn.Linear(
-            self.hidden_size, self.hidden_size * 3, bias=False, dtype=torch.float16
-        )
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            rotary_dim=self.head_dim,
-            max_position=max_position,
-            base=base,
-        )
-
-    def forward(self, positions, hidden_states):
-        # Simulate the pattern: mm -> split_with_sizes -> rotary_embedding
-        # -> slice_scatter -> split_with_sizes
-
-        qkv = self.qkv_proj(hidden_states)
-        split_sizes = [self.hidden_size, self.hidden_size, self.hidden_size]
-        q, k, v = torch.split(qkv, split_sizes, dim=-1)
-
-        q_rotated, k_rotated = self.rotary_emb(positions, q, k)
-
-        qkv_updated = torch.cat([q_rotated, k_rotated, v], dim=-1)
-        return qkv_updated
-
-    def example_inputs(self, num_tokens=32, head_dim=64, num_heads=4):
-        dtype = torch.float16
-        hidden_size = head_dim * num_heads
-        positions = torch.arange(num_tokens, dtype=torch.long)
-        hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
-        return (positions, hidden_states)
-
-    def ops_in_model(self, do_fusion):
-        return [torch.ops._C.rotary_embedding.default]
-
-    def ops_not_in_model(self):
-        return [torch.ops.aten.slice_scatter.default]
-
-
-MODELS = [
-    TestSiluMul,
-    TestFusedAddRMSNorm,
-    TestRotaryEmbedding,
-    TestRotaryEmbeddingSliceScatter,
+SILU_MUL_QUANT_OP = torch.ops._C.silu_and_mul_quant.default
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
 ]
 
 
-@pytest.mark.parametrize("model_class", MODELS)
+@pytest.mark.parametrize(
+    "model, quant_key",
+    [("nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e", kFp8StaticTensorSym),
+     ("nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8_DYNAMIC-e2e",
+      kFp8DynamicTokenSym)])
 @pytest.mark.parametrize("do_fusion", [True, False])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
-def test_fix_functionalization(model_class: torch.nn.Module, do_fusion: bool):
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda",
+                    reason="Only test on CUDA")
+def test_fix_functionalization(model: str, quant_key: QuantKey,
+                               do_fusion: bool):
     torch.set_default_device("cuda")
 
     vllm_config = VllmConfig()
     vllm_config.compilation_config = CompilationConfig(
-        pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True)
-    )
+        pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True))
     noop_pass = NoOpEliminationPass(vllm_config)
     fusion_pass = RMSNormQuantFusionPass(vllm_config)
     cleanup_pass = PostCleanupPass(vllm_config)
     act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)
 
-    passes = (
-        [noop_pass, fusion_pass, act_quant_fusion_pass, cleanup_pass]
-        if do_fusion
-        else [noop_pass, cleanup_pass]
-    )
+    passes = [noop_pass, fusion_pass, act_quant_fusion_pass, cleanup_pass
+              ] if do_fusion else [noop_pass, cleanup_pass]
     func_pass = FixFunctionalizationPass(vllm_config)
-
     backend_func = TestBackend(*passes, func_pass)
     backend_no_func = TestBackend(*passes)
 
-    model = model_class()
-    torch.compile(model, backend=backend_func)(*model.example_inputs())
-    torch.compile(model, backend=backend_no_func)(*model.example_inputs())
+    # instantiate a full engine and manually compile the model 2x
+    # (with and without FixFunctionalizationPass)
+    llm = LLM(model=model, enforce_eager=True)
+    model_runner = llm.llm_engine.model_executor.driver_worker.model_runner
+    orig_model = model_runner.model
+    # TODO mark inputs dynamic? (currently torch.compile is triggered 4x)
+    # Can only do that by using the decorator but then we'd have to instantiate
+    # 2 LLM instances.
 
-    # check if the functionalization pass is applied
-    for op in model.ops_in_model(do_fusion):
+    sampling_params = SamplingParams(temperature=0.0, top_p=1.0)
+    model_runner.model = torch.compile(orig_model,
+                                       fullgraph=True,
+                                       backend=backend_func)
+    gen_func = llm.generate(prompts, sampling_params)
+
+    model_runner.model = torch.compile(orig_model,
+                                       fullgraph=True,
+                                       backend=backend_no_func)
+
+    gen_no_func = llm.generate(prompts, sampling_params)
+
+    for output_func, output_no_func in zip(gen_func, gen_no_func):
+        assert output_func.outputs[0].text == output_no_func.outputs[0].text
+
+    # OPS_IN_MODEL always appear. RMS_OP is fused away if we run fusion,
+    # and replaced by fused quantized ops in RMS_QUANT_OPS.
+    rms_ops = [FUSED_OPS[(quant_key, True)], FUSED_OPS[(quant_key, False)]
+               ] if do_fusion else [RMS_OP]
+    silu_mul_ops = [SILU_MUL_QUANT_OP] if do_fusion and \
+        quant_key == kFp8StaticTensorSym else [
+        SILU_MUL_OP
+    ]
+
+    ops = OPS_IN_MODEL + rms_ops + silu_mul_ops
+
+    for op in ops:
         find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
-        assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None
+        assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes,
+                                  op) is None  # noqa: E501
 
     # make sure the ops were all de-functionalized
     found = dict()
     for node in backend_func.graph_post_pass.nodes:
-        for op in model.ops_in_model(do_fusion):
+        for op in ops:
             if is_func(node, op):
                 found[op] = True
-        for op in model.ops_not_in_model():
-            if is_func(node, op):
-                found[op] = True
-    assert all(found[op] for op in model.ops_in_model(do_fusion))
-    assert all(not found.get(op) for op in model.ops_not_in_model())
+    assert all(found[op] for op in ops)

tests/compile/test_fusion.py

@@ -5,26 +5,17 @@ import pytest
 import torch
 
 import vllm.plugins
-from vllm.compilation.fusion import (
-    FUSED_OPS,
-    QUANT_OPS,
-    FusedRMSQuantKey,
-    RMSNormQuantFusionPass,
-)
+from vllm.compilation.fusion import (FUSED_OPS, QUANT_OPS, FusedRMSQuantKey,
+                                     RMSNormQuantFusionPass)
 from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
-from vllm.config import CompilationConfig, CompilationLevel, PassConfig, VllmConfig
+from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,
+                         VllmConfig)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape,
-    QuantKey,
-    ScaleDesc,
-)
+    GroupShape, QuantKey, ScaleDesc)
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
-    Fp8LinearOp,
-    cutlass_fp8_supported,
-    maybe_create_device_identity,
-)
+    Fp8LinearOp, cutlass_fp8_supported, maybe_create_device_identity)
 from vllm.platforms import current_platform
 
 from ..utils import override_cutlass_fp8_supported
@@ -34,15 +25,9 @@ FP8_DTYPE = current_platform.fp8_dtype()
 
 
 class TestModel(torch.nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        eps: float,
-        static: bool,
-        cuda_force_torch: bool,
-        *args,
-        **kwargs,
-    ):
+
+    def __init__(self, hidden_size: int, eps: float, static: bool,
+                 cuda_force_torch: bool, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.cuda_force_torch = cuda_force_torch
         self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]
@@ -69,15 +54,17 @@ class TestModel(torch.nn.Module):
         resid = torch.sqrt(x)
         y = self.norm[0](x)
 
-        x2 = self.fp8_linear.apply(
-            y, self.w[0], self.wscale[0], input_scale=self.scale[0]
-        )
+        x2 = self.fp8_linear.apply(y,
+                                   self.w[0],
+                                   self.wscale[0],
+                                   input_scale=self.scale[0])
         # make sure resid is used for replacement to work
         y2, resid = self.norm[1](x2, resid)
 
-        x3 = self.fp8_linear.apply(
-            y2, self.w[1], self.wscale[1], input_scale=self.scale[1]
-        )
+        x3 = self.fp8_linear.apply(y2,
+                                   self.w[1],
+                                   self.wscale[1],
+                                   input_scale=self.scale[1])
         y3, resid = self.norm[2](x3, resid)  # use resid here
         return y3
 
@@ -87,7 +74,7 @@ class TestModel(torch.nn.Module):
     def ops_in_model_after(self):
         return [
             FUSED_OPS[FusedRMSQuantKey(self.key, False)],
-            FUSED_OPS[FusedRMSQuantKey(self.key, True)],
+            FUSED_OPS[FusedRMSQuantKey(self.key, True)]
         ]
 
 
@@ -98,27 +85,22 @@ class TestModel(torch.nn.Module):
 @pytest.mark.parametrize("static", [True, False])
 # cuda_force_torch used to test torch code path on platforms that
 # cutlass_fp8_supported() == True.
-@pytest.mark.parametrize(
-    "cuda_force_torch", [True, False] if cutlass_fp8_supported() else [True]
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(), reason="Only test on CUDA and ROCm"
-)
-def test_fusion_rmsnorm_quant(
-    dtype, hidden_size, num_tokens, eps, static, cuda_force_torch
-):
+@pytest.mark.parametrize("cuda_force_torch",
+                         [True, False] if cutlass_fp8_supported() else [True])
+@pytest.mark.skipif(not current_platform.is_cuda_alike(),
+                    reason="Only test on CUDA and ROCm")
+def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
+                              cuda_force_torch):
     torch.set_default_device("cuda")
     torch.set_default_dtype(dtype)
     torch.manual_seed(1)
     maybe_create_device_identity()  # needed for certain non-cutlass fp8 paths
 
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            custom_ops=["+rms_norm", "+quant_fp8"],
-            pass_config=PassConfig(enable_fusion=True, enable_noop=True),
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        custom_ops=["+rms_norm", "+quant_fp8"],
+        pass_config=PassConfig(enable_fusion=True, enable_noop=True),
+    ))
     with vllm.config.set_current_vllm_config(vllm_config):
         # Reshape pass is needed for the fusion pass to work
         noop_pass = NoOpEliminationPass(vllm_config)

tests/compile/test_fusion_all_reduce.py

@@ -10,24 +10,14 @@ from vllm.compilation.collective_fusion import AllReduceFusionPass
 from vllm.compilation.fix_functionalization import FixFunctionalizationPass
 from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
-from vllm.config import (
-    CompilationConfig,
-    CompilationLevel,
-    DeviceConfig,
-    ModelConfig,
-    PassConfig,
-    VllmConfig,
-)
+from vllm.config import (CompilationConfig, CompilationLevel, DeviceConfig,
+                         ModelConfig, PassConfig, VllmConfig)
 from vllm.distributed import tensor_model_parallel_all_reduce
-from vllm.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-)
+from vllm.distributed.parallel_state import (init_distributed_environment,
+                                             initialize_model_parallel)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
-    GroupShape,
-    QuantFP8,
-)
+    GroupShape, QuantFP8)
 from vllm.platforms import current_platform
 from vllm.utils import update_environment_variables
 
@@ -36,6 +26,7 @@ from .backend import TestBackend
 
 
 class TestAllReduceRMSNormModel(torch.nn.Module):
+
     def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
         super().__init__()
         self.hidden_size = hidden_size
@@ -56,6 +47,7 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
 
 
 class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):
+
     def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
         super().__init__()
         self.hidden_size = hidden_size
@@ -76,22 +68,25 @@ class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):
 
 
 class TestAllReduceFusedAddRMSNormStaticQuantFP8Model(torch.nn.Module):
+
     def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
         super().__init__()
         self.hidden_size = hidden_size
         self.eps = eps
         self.norm = RMSNorm(hidden_size, eps)
-        self.quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
+        self.quant_fp8 = QuantFP8(static=True,
+                                  group_shape=GroupShape.PER_TENSOR)
         self.scale = torch.rand(1, dtype=torch.float32)
-        self.output = torch.empty((token_num, hidden_size), dtype=torch.float32)
+        self.output = torch.empty((token_num, hidden_size),
+                                  dtype=torch.float32)
 
     def forward(self, hidden_states, residual):
         view = hidden_states.reshape(-1, self.hidden_size)
         all_reduce = tensor_model_parallel_all_reduce(view)
         norm_output, residual_output = self.norm(all_reduce, residual)
-        torch.ops._C.static_scaled_fp8_quant(
-            self.output, norm_output.contiguous(), self.scale
-        )
+        torch.ops._C.static_scaled_fp8_quant(self.output,
+                                             norm_output.contiguous(),
+                                             self.scale)
         return self.output, residual_output
 
     def ops_in_model_after(self):
@@ -100,33 +95,35 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP8Model(torch.nn.Module):
     def ops_in_model_before(self):
         return [
             torch.ops.vllm.all_reduce.default,
-            torch.ops._C.static_scaled_fp8_quant.default,
+            torch.ops._C.static_scaled_fp8_quant.default
         ]
 
 
 class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
+
     def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
         super().__init__()
         self.hidden_size = hidden_size
         self.eps = eps
         self.norm = RMSNorm(hidden_size, eps)
         self.scale = torch.rand(1, dtype=torch.float32)
-        self.output = torch.empty((token_num, hidden_size), dtype=torch.float32)
+        self.output = torch.empty((token_num, hidden_size),
+                                  dtype=torch.float32)
 
         round_up = lambda x, y: (x + y - 1) // y * y
         rounded_m = round_up(token_num, 128)
         scale_n = hidden_size // 16
         rounded_n = round_up(scale_n, 4)
-        self.output_scale = torch.empty((rounded_m, rounded_n // 4), dtype=torch.int32)
+        self.output_scale = torch.empty((rounded_m, rounded_n // 4),
+                                        dtype=torch.int32)
 
     def forward(self, hidden_states, residual):
         view = hidden_states.reshape(-1, self.hidden_size)
         all_reduce = tensor_model_parallel_all_reduce(view)
         norm_output, residual_output = self.norm(all_reduce, residual)
         norm_output = norm_output.reshape(-1, norm_output.shape[-1])
-        torch.ops._C.scaled_fp4_quant(
-            self.output, norm_output, self.output_scale, self.scale
-        )
+        torch.ops._C.scaled_fp4_quant(self.output, norm_output,
+                                      self.output_scale, self.scale)
         return self.output, residual_output, self.output_scale
 
     def ops_in_model_after(self):
@@ -135,7 +132,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
     def ops_in_model_before(self):
         return [
             torch.ops.vllm.all_reduce.default,
-            torch.ops._C.scaled_fp4_quant.default,
+            torch.ops._C.scaled_fp4_quant.default
         ]
 
 
@@ -148,55 +145,41 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
         TestAllReduceFusedAddRMSNormStaticQuantFP8Model,
         # TODO: Enable with torch==2.8.0
         # TestAllReduceFusedAddRMSNormStaticQuantFP4Model,
-    ],
-)
+    ])
 @pytest.mark.parametrize("batch_size", [8])
 @pytest.mark.parametrize("seq_len", [8])
 @pytest.mark.parametrize("hidden_size", [16])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
+                    reason="Only test on CUDA")
 @pytest.mark.skipif(
     not find_spec("flashinfer")
     or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),
     reason="flashinfer is not found or flashinfer "
-    "is not compiled with trtllm_allreduce_fusion",
-)
-def test_all_reduce_fusion_pass_replace(
-    test_model: torch.nn.Module,
-    batch_size: int,
-    seq_len: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-):
+    "is not compiled with trtllm_allreduce_fusion")
+def test_all_reduce_fusion_pass_replace(test_model: torch.nn.Module,
+                                        batch_size: int, seq_len: int,
+                                        hidden_size: int, dtype: torch.dtype):
     num_processes = 2
-    if (
-        test_model == TestAllReduceFusedAddRMSNormStaticQuantFP4Model
-        and not current_platform.has_device_capability(100)
-    ):
-        pytest.skip(
-            "Skip as nvfp4 is only supported on "
-            "devices with compute capability 10.0 (Blackwell)"
-        )
+    if (test_model == TestAllReduceFusedAddRMSNormStaticQuantFP4Model
+            and not current_platform.has_device_capability(100)):
+        pytest.skip("Skip as nvfp4 is only supported on "
+                    "devices with compute capability 10.0 (Blackwell)")
 
     def run_torch_spawn(fn, nprocs):
-        torch.multiprocessing.spawn(
-            fn,
-            args=(num_processes, test_model, batch_size, seq_len, hidden_size, dtype),
-            nprocs=nprocs,
-        )
+        torch.multiprocessing.spawn(fn,
+                                    args=(num_processes, test_model,
+                                          batch_size, seq_len, hidden_size,
+                                          dtype),
+                                    nprocs=nprocs)
 
     run_torch_spawn(all_reduce_fusion_pass_on_test_model, num_processes)
 
 
-def all_reduce_fusion_pass_on_test_model(
-    local_rank: int,
-    world_size: int,
-    test_model_cls: torch.nn.Module,
-    batch_size: int,
-    seq_len: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-):
+def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
+                                         test_model_cls: torch.nn.Module,
+                                         batch_size: int, seq_len: int,
+                                         hidden_size: int, dtype: torch.dtype):
     current_platform.seed_everything(0)
 
     device = torch.device(f"cuda:{local_rank}")
@@ -204,42 +187,39 @@ def all_reduce_fusion_pass_on_test_model(
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "12345",
-        }
-    )
+    update_environment_variables({
+        'RANK': str(local_rank),
+        'LOCAL_RANK': str(local_rank),
+        'WORLD_SIZE': str(world_size),
+        'MASTER_ADDR': 'localhost',
+        'MASTER_PORT': '12345',
+    })
 
     init_distributed_environment()
     initialize_model_parallel(tensor_model_parallel_size=world_size)
 
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm", "+quant_fp8"]
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        custom_ops=["+rms_norm", "+quant_fp8"]))
     vllm_config.compilation_config.pass_config = PassConfig(
-        enable_fi_allreduce_fusion=True, enable_noop=True
-    )
+        enable_fi_allreduce_fusion=True, enable_noop=True)
     vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.
     model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
-    vllm_config.model_config = ModelConfig(
-        model=model_name, trust_remote_code=True, dtype=dtype, seed=42
-    )
+    vllm_config.model_config = ModelConfig(model=model_name,
+                                           trust_remote_code=True,
+                                           dtype=dtype,
+                                           seed=42)
 
     all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)
     noop_pass = NoOpEliminationPass(vllm_config)
     func_pass = FixFunctionalizationPass(vllm_config)
     cleanup_pass = PostCleanupPass(vllm_config)
 
-    backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass, cleanup_pass)
+    backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass,
+                          cleanup_pass)
 
     token_num = batch_size * seq_len
     model = test_model_cls(hidden_size, token_num)

tests/compile/test_fusion_attn.py

@@ -19,23 +19,14 @@ from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
 from vllm.compilation.fx_utils import find_op_nodes
 from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    CompilationLevel,
-    ModelConfig,
-    PassConfig,
-    SchedulerConfig,
-    VllmConfig,
-    set_current_vllm_config,
-)
+from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
+                         ModelConfig, PassConfig, SchedulerConfig, VllmConfig,
+                         set_current_vllm_config)
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    QuantKey,
-    kFp8StaticTensorSym,
-    kNvfp4Quant,
-)
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
+    QuantKey, kFp8StaticTensorSym, kNvfp4Quant)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    Fp8LinearOp)
 from vllm.platforms import current_platform
 from vllm.utils import is_torch_equal_or_newer
 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -49,16 +40,14 @@ backend_unfused: Optional[TestBackend] = None
 
 
 @pytest.mark.parametrize(
-    "model, quant_key", [("amd/Llama-3.1-8B-Instruct-FP8-KV", kFp8StaticTensorSym)]
-)
+    "model, quant_key",
+    [("amd/Llama-3.1-8B-Instruct-FP8-KV", kFp8StaticTensorSym)])
 @pytest.mark.parametrize("use_triton_fa", [True, False])
 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
-@pytest.mark.skipif(
-    not current_platform.is_rocm(), reason="V0 attn quant fusion only on ROCm"
-)
-def test_attention_fusion_v0(
-    example_prompts, monkeypatch, model: str, quant_key: QuantKey, use_triton_fa: bool
-):
+@pytest.mark.skipif(not current_platform.is_rocm(),
+                    reason="V0 attn quant fusion only on ROCm")
+def test_attention_fusion_v0(example_prompts, monkeypatch, model: str,
+                             quant_key: QuantKey, use_triton_fa: bool):
     # Clean Dynamo cache to avoid reusing other test cases
     # (for some reason the reset at the end is not enough)
     torch._dynamo.reset()
@@ -80,24 +69,22 @@ def test_attention_fusion_v0(
         backend="tests.compile.test_fusion_attn.backend_unfused",
         custom_ops=["+quant_fp8"],
     )
-    vllm_config = VllmConfig(
-        compilation_config=compile_config,
-        model_config=ModelConfig(
-            model=model,
-            dtype=torch.bfloat16,
-        ),
-    )
+    vllm_config = VllmConfig(compilation_config=compile_config,
+                             model_config=ModelConfig(
+                                 model=model,
+                                 dtype=torch.bfloat16,
+                             ))
     backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
 
-    llm = LLM(
-        model,
-        enforce_eager=True,
-        compilation_config=compile_config,
-        gpu_memory_utilization=0.5,
-        max_model_len=2048,
-    )
+    llm = LLM(model,
+              enforce_eager=True,
+              compilation_config=compile_config,
+              gpu_memory_utilization=0.5,
+              max_model_len=2048)
 
-    sampling_params = SamplingParams(temperature=0.0, max_tokens=10, top_p=0.95)
+    sampling_params = SamplingParams(temperature=0.0,
+                                     max_tokens=10,
+                                     top_p=0.95)
 
     unfused_output = llm.generate(prompts, sampling_params)
     backend_unfused = None  # Reset backend to make sure llm gets released
@@ -110,25 +97,21 @@ def test_attention_fusion_v0(
         backend="tests.compile.test_fusion_attn.backend",
         custom_ops=["+quant_fp8"],
     )
-    vllm_config = VllmConfig(
-        compilation_config=compile_config,
-        model_config=ModelConfig(
-            model=model,
-            dtype=torch.bfloat16,
-        ),
-    )
+    vllm_config = VllmConfig(compilation_config=compile_config,
+                             model_config=ModelConfig(
+                                 model=model,
+                                 dtype=torch.bfloat16,
+                             ))
 
     # AttnFusionPass needs attention layers to be registered in config upon init
     # so we initialize it during compilation.
     attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
     backend = TestBackend(NoOpEliminationPass(vllm_config), attn_pass)
-    llm2 = LLM(
-        model,
-        enforce_eager=True,
-        compilation_config=compile_config,
-        gpu_memory_utilization=0.5,
-        max_model_len=2048,
-    )
+    llm2 = LLM(model,
+               enforce_eager=True,
+               compilation_config=compile_config,
+               gpu_memory_utilization=0.5,
+               max_model_len=2048)
 
     # check support
     attn_fusion_supported = [
@@ -149,9 +132,9 @@ def test_attention_fusion_v0(
     for i in range(len(attn_nodes_pre)):
         assert attn_nodes_pre[i].kwargs["output_scale"] is None
         fused = attn_nodes_post[i].kwargs["output_scale"] is not None
-        assert fused == attn_fusion_supported[i], (
-            f"Node {i} {'' if fused else 'not '} expected to have fused output quant"
-        )
+        assert fused == attn_fusion_supported[i], \
+            f"Node {i} {'' if fused else 'not '} expected " \
+            f"to have fused output quant"
 
     # check outputs
     fused_output = llm2.generate(prompts, sampling_params)
@@ -177,16 +160,9 @@ def test_attention_fusion_v0(
 class AttentionQuantPatternModel(torch.nn.Module):
     """Base model for AttentionQuantPattern fusion."""
 
-    def __init__(
-        self,
-        num_qo_heads: int,
-        num_kv_heads: int,
-        head_size: int,
-        kv_cache_dtype: torch.dtype,
-        device: torch.device,
-        vllm_config: VllmConfig,
-        **kwargs,
-    ):
+    def __init__(self, num_qo_heads: int, num_kv_heads: int, head_size: int,
+                 kv_cache_dtype: torch.dtype, device: torch.device,
+                 vllm_config: VllmConfig, **kwargs):
         super().__init__()
         self.num_qo_heads = num_qo_heads
         self.num_kv_heads = num_kv_heads
@@ -221,30 +197,33 @@ class AttentionQuantPatternModel(torch.nn.Module):
             device=self.device,
         )
 
-    def build_attn_metadata(self, batch_size: int, use_hnd: bool) -> AttentionMetadata:
+    def build_attn_metadata(self, batch_size: int, use_hnd: bool) \
+            -> AttentionMetadata:
         """Initialize attention metadata."""
 
         # Create common attn metadata
-        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
+        batch_spec = BatchSpec(seq_lens=[1] * batch_size,
+                               query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
-            batch_spec, self.block_size, self.device, arange_block_indices=True
-        )
+            batch_spec,
+            self.block_size,
+            self.device,
+            arange_block_indices=True)
 
-        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
+        max_blocks = (max(batch_spec.seq_lens) + self.block_size -
+                      1) // self.block_size
         num_blocks = batch_size * max_blocks
 
         # Create dummy KV cache for FlashInfer TRTLLM
         #   - NHD: [num_blocks, block_size, num_kv_heads, head_size]
         #   - HND: [num_blocks, num_kv_heads, block_size, head_size]
-        kv_cache = torch.zeros(
-            num_blocks,
-            2,
-            self.num_kv_heads,
-            self.block_size,
-            self.head_size,
-            dtype=self.kv_cache_dtype,
-            device=self.device,
-        )
+        kv_cache = torch.zeros(num_blocks,
+                               2,
+                               self.num_kv_heads,
+                               self.block_size,
+                               self.head_size,
+                               dtype=self.kv_cache_dtype,
+                               device=self.device)
         if current_platform.is_rocm():
             # k/v as 1st dimention
             if use_hnd:
@@ -260,8 +239,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
 
         # Build attn metadata
         self.attn_metadata = self.builder.build(
-            common_prefix_len=0, common_attn_metadata=common_attn_metadata
-        )
+            common_prefix_len=0, common_attn_metadata=common_attn_metadata)
 
         return self.attn_metadata
 
@@ -276,30 +254,27 @@ class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
 
         self.fp8_linear = Fp8LinearOp(
             act_quant_static=self.quant_key.scale.static,
-            act_quant_group_shape=self.quant_key.scale.group_shape,
-        )
+            act_quant_group_shape=self.quant_key.scale.group_shape)
 
         hidden_size = self.num_qo_heads * self.head_size
         self.w = kwargs.get(
-            "w",
-            {
-                "weight": torch.randn(hidden_size, hidden_size)
-                .to(dtype=FP8_DTYPE, device=self.device)
-                .t(),
-                "wscale": torch.tensor([1.0], dtype=torch.float32, device=self.device),
-                "scale": torch.tensor([1.0], dtype=torch.float32, device=self.device),
-            },
-        )
+            "w", {
+                "weight":
+                torch.randn(hidden_size, hidden_size).to(
+                    dtype=FP8_DTYPE, device=self.device).t(),
+                "wscale":
+                torch.tensor([1.0], dtype=torch.float32, device=self.device),
+                "scale":
+                torch.tensor([1.0], dtype=torch.float32, device=self.device),
+            })
 
     def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
         """Forward pass that creates the pattern to be fused."""
         attn_output = self.attn(q, k, v)
-        return self.fp8_linear.apply(
-            input=attn_output,
-            weight=self.w["weight"],
-            weight_scale=self.w["wscale"],
-            input_scale=self.w["scale"],
-        )
+        return self.fp8_linear.apply(input=attn_output,
+                                     weight=self.w["weight"],
+                                     weight_scale=self.w["wscale"],
+                                     input_scale=self.w["scale"])
 
 
 class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
@@ -312,54 +287,42 @@ class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
 
         hidden_size = self.num_qo_heads * self.head_size
         self.w = kwargs.get(
-            "w",
-            {
-                "weight": torch.randint(
-                    256,
-                    (hidden_size, hidden_size // 2),
-                    dtype=FP4_DTYPE,
-                    device=self.device,
-                ),
-                "wscale_swizzled": torch.randn(hidden_size, hidden_size // 16).to(
-                    dtype=FP8_DTYPE, device=self.device
-                ),
-                "wscale": torch.tensor([500], dtype=torch.float32, device=self.device),
-                "scale": torch.tensor([0.002], dtype=torch.float32, device=self.device),
-            },
-        )
+            "w", {
+                "weight":
+                torch.randint(256, (hidden_size, hidden_size // 2),
+                              dtype=FP4_DTYPE,
+                              device=self.device),
+                "wscale_swizzled":
+                torch.randn(hidden_size, hidden_size // 16).to(
+                    dtype=FP8_DTYPE, device=self.device),
+                "wscale":
+                torch.tensor([500], dtype=torch.float32, device=self.device),
+                "scale":
+                torch.tensor([0.002], dtype=torch.float32, device=self.device),
+            })
 
     def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
         """Forward pass that creates the pattern to be fused."""
         attn_output = self.attn(q, k, v)
         quant_output, output_block_scale = scaled_fp4_quant(
-            attn_output, 1 / self.w["scale"]
-        )
-        return cutlass_scaled_fp4_mm(
-            a=quant_output,
-            b=self.w["weight"],
-            block_scale_a=output_block_scale,
-            block_scale_b=self.w["wscale_swizzled"],
-            alpha=self.w["scale"] * self.w["wscale"],
-            out_dtype=attn_output.dtype,
-        )
+            attn_output, 1 / self.w["scale"])
+        return cutlass_scaled_fp4_mm(a=quant_output,
+                                     b=self.w["weight"],
+                                     block_scale_a=output_block_scale,
+                                     block_scale_b=self.w["wscale_swizzled"],
+                                     alpha=self.w["scale"] * self.w["wscale"],
+                                     out_dtype=attn_output.dtype)
 
 
 if current_platform.is_cuda():
-    MODELS = [
-        (
-            "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
-            TestAttentionFp8StaticQuantPatternModel,
-        ),
-        (
-            "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
-            TestAttentionNvfp4QuantPatternModel,
-        ),
-    ]
+    MODELS = [("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
+               TestAttentionFp8StaticQuantPatternModel),
+              ("nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
+               TestAttentionNvfp4QuantPatternModel)]
     HEADS = [(64, 8), (40, 8)]
 elif current_platform.is_rocm():
-    MODELS = [
-        ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)
-    ]
+    MODELS = [("amd/Llama-3.1-8B-Instruct-FP8-KV",
+               TestAttentionFp8StaticQuantPatternModel)]
     HEADS = [(32, 8), (40, 8)]
 else:
     MODELS = []
@@ -368,53 +331,41 @@ else:
 
 @pytest.mark.parametrize("num_qo_heads, num_kv_heads", HEADS)
 @pytest.mark.parametrize("head_size", [128])
-@pytest.mark.parametrize(
-    "batch_size", [7, 256, 533] if current_platform.is_cuda() else [8]
-)
+@pytest.mark.parametrize("batch_size",
+                         [7, 256, 533] if current_platform.is_cuda() else [8])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("model_name, model_class", MODELS)
+@pytest.mark.parametrize("backend",
+                         [_Backend.FLASHINFER] if current_platform.is_cuda()
+                         else [_Backend.TRITON_ATTN])
 @pytest.mark.parametrize(
-    "backend",
-    [_Backend.FLASHINFER] if current_platform.is_cuda() else [_Backend.TRITON_ATTN],
-)
-@pytest.mark.parametrize(
-    "split_attention", [False, True] if current_platform.is_rocm() else [False]
-)
+    "split_attention",
+    [False, True] if current_platform.is_rocm() else [False])
 # TODO(boyuan): test inductor graph partition on rocm
 @pytest.mark.parametrize(
     "use_inductor_graph_partition",
-    [False] if current_platform.is_rocm() else [False, True],
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"
-)
+    [False] if current_platform.is_rocm() else [False, True])
+@pytest.mark.skipif(not current_platform.is_cuda_alike(),
+                    reason="Only test ROCm or CUDA")
 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
-@pytest.mark.skipif(
-    current_platform.is_cuda() and not current_platform.is_device_capability((10, 0)),
-    reason="On CUDA only test on SM100(Blackwell)",
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"
-)
-def test_attention_quant_pattern(
-    num_qo_heads: int,
-    num_kv_heads: int,
-    head_size: int,
-    batch_size: int,
-    dtype: torch.dtype,
-    model_name: str,
-    model_class: type[AttentionQuantPatternModel],
-    backend: _Backend,
-    split_attention: bool,
-    use_inductor_graph_partition: bool,
-    monkeypatch,
-    dist_init,
-    caplog_vllm,
-):
+@pytest.mark.skipif(current_platform.is_cuda()
+                    and not current_platform.is_device_capability((10, 0)),
+                    reason="On CUDA only test on SM100(Blackwell)")
+@pytest.mark.skipif(not current_platform.is_cuda_alike(),
+                    reason="Only test ROCm or CUDA")
+def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
+                                 head_size: int, batch_size: int,
+                                 dtype: torch.dtype, model_name: str,
+                                 model_class: type[AttentionQuantPatternModel],
+                                 backend: _Backend, split_attention: bool,
+                                 use_inductor_graph_partition: bool,
+                                 monkeypatch, dist_init, caplog_vllm):
     """Test AttentionStaticQuantPattern fusion pass"""
 
-    if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
-        pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
+    if use_inductor_graph_partition and not is_torch_equal_or_newer(
+            "2.9.0.dev"):
+        pytest.skip("inductor graph partition is only available "
+                    "in PyTorch 2.9+")
 
     monkeypatch.setenv("VLLM_USE_V1", "1")
     if split_attention:
@@ -435,13 +386,21 @@ def test_attention_quant_pattern(
             custom_ops=["+quant_fp8"],
             use_inductor_graph_partition=use_inductor_graph_partition,
         ),
-        cache_config=CacheConfig(cache_dtype="fp8"),
-    )
+        cache_config=CacheConfig(cache_dtype="fp8"))
 
     # Create test inputs
-    q = torch.randn(batch_size, num_qo_heads * head_size, dtype=dtype, device=device)
-    k = torch.randn(batch_size, num_kv_heads * head_size, dtype=dtype, device=device)
-    v = torch.randn(batch_size, num_kv_heads * head_size, dtype=dtype, device=device)
+    q = torch.randn(batch_size,
+                    num_qo_heads * head_size,
+                    dtype=dtype,
+                    device=device)
+    k = torch.randn(batch_size,
+                    num_kv_heads * head_size,
+                    dtype=dtype,
+                    device=device)
+    v = torch.randn(batch_size,
+                    num_kv_heads * head_size,
+                    dtype=dtype,
+                    device=device)
 
     # Mark first dimension as dynamic for realistic testing
     torch._dynamo.mark_dynamic(q, 0)
@@ -450,53 +409,42 @@ def test_attention_quant_pattern(
 
     # Run model directly without compilation and fusion
     vllm_config_unfused = copy.deepcopy(vllm_config)
-    with (
-        set_current_vllm_config(vllm_config_unfused),
-        set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused),
-        global_force_attn_backend_context_manager(backend),
-    ):
-        model_unfused = model_class(
-            num_qo_heads=num_qo_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            kv_cache_dtype=FP8_DTYPE,
-            device=device,
-            vllm_config=vllm_config_unfused,
-        )
+    with set_current_vllm_config(vllm_config_unfused), set_forward_context(
+            attn_metadata=None, vllm_config=vllm_config_unfused
+    ), global_force_attn_backend_context_manager(backend):
+        model_unfused = model_class(num_qo_heads=num_qo_heads,
+                                    num_kv_heads=num_kv_heads,
+                                    head_size=head_size,
+                                    kv_cache_dtype=FP8_DTYPE,
+                                    device=device,
+                                    vllm_config=vllm_config_unfused)
         model_unfused = model_unfused.to(device)
 
         forward_ctx = get_forward_context()
         forward_ctx.attn_metadata = model_unfused.build_attn_metadata(
-            batch_size, use_hnd=split_attention
-        )
+            batch_size, use_hnd=split_attention)
 
         # Run model directly without compilation and fusion
         result_unfused = model_unfused(q, k, v)
 
     # Run model with attn fusion enabled
     vllm_config.compilation_config.pass_config = PassConfig(
-        enable_attn_fusion=True, enable_noop=True
-    )
-    with (
-        set_current_vllm_config(vllm_config),
-        set_forward_context(attn_metadata=None, vllm_config=vllm_config),
-        global_force_attn_backend_context_manager(backend),
-    ):
-        model_fused = model_class(
-            num_qo_heads=num_qo_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            kv_cache_dtype=FP8_DTYPE,
-            device=device,
-            vllm_config=vllm_config,
-            w=model_unfused.w,
-        )
+        enable_attn_fusion=True, enable_noop=True)
+    with set_current_vllm_config(vllm_config), set_forward_context(
+            attn_metadata=None, vllm_config=vllm_config
+    ), global_force_attn_backend_context_manager(backend):
+        model_fused = model_class(num_qo_heads=num_qo_heads,
+                                  num_kv_heads=num_kv_heads,
+                                  head_size=head_size,
+                                  kv_cache_dtype=FP8_DTYPE,
+                                  device=device,
+                                  vllm_config=vllm_config,
+                                  w=model_unfused.w)
         model_fused = model_fused.to(device)
 
         forward_ctx = get_forward_context()
         forward_ctx.attn_metadata = model_fused.build_attn_metadata(
-            batch_size, use_hnd=split_attention
-        )
+            batch_size, use_hnd=split_attention)
 
         # Create test backend with fusion passes enabled
         noop_pass = NoOpEliminationPass(vllm_config)
@@ -506,9 +454,9 @@ def test_attention_quant_pattern(
         test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
 
         # Compile model with fusion enabled
-        model_compiled = torch.compile(
-            model_fused, backend=test_backend, fullgraph=True
-        )
+        model_compiled = torch.compile(model_fused,
+                                       backend=test_backend,
+                                       fullgraph=True)
         assert model_compiled.attn._o_scale_float is None
 
         result_fused_1 = model_compiled(q, k, v)
@@ -523,49 +471,49 @@ def test_attention_quant_pattern(
 
             assert model_compiled.attn._o_scale_float is not None
 
-            torch.testing.assert_close(
-                result_unfused, result_fused_2, atol=1e-2, rtol=1e-2
-            )
+            torch.testing.assert_close(result_unfused,
+                                       result_fused_2,
+                                       atol=1e-2,
+                                       rtol=1e-2)
 
     # Check attn fusion support
     quant_key = model_class.quant_key
     attn_fusion_supported = [
-        layer.impl.fused_output_quant_supported(quant_key)
-        for key, layer in vllm_config.compilation_config.static_forward_context.items()
+        layer.impl.fused_output_quant_supported(quant_key) for key, layer in
+        vllm_config.compilation_config.static_forward_context.items()
     ]
     if any(attn_fusion_supported):
         # Check quantization ops in the graph before and after fusion
-        test_backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=True)
+        test_backend.check_before_ops([QUANT_OPS[quant_key]],
+                                      fully_replaced=True)
 
     # access the underlying `AttnFusionPass` on the `LazyInitPass`
     assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
 
     # Check attention ops in the graph before and after fusion
     attn_nodes_pre = list(find_op_nodes(ATTN_OP, test_backend.graph_pre_pass))
-    attn_nodes_post = list(find_op_nodes(ATTN_OP, test_backend.graph_post_pass))
+    attn_nodes_post = list(find_op_nodes(ATTN_OP,
+                                         test_backend.graph_post_pass))
 
     assert len(attn_nodes_pre) > 0, "Should have attention nodes before fusion"
-    assert len(attn_nodes_pre) == len(attn_nodes_post), (
+    assert len(attn_nodes_pre) == len(attn_nodes_post), \
         "Should have same number of attention nodes before and after fusion"
-    )
-    assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
+    assert attn_nodes_pre[0].kwargs.get("output_scale") is None, \
         "Attention should not have output_scale before fusion"
-    )
-    assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
+    assert attn_nodes_post[0].kwargs.get("output_scale") is not None, \
         "Attention should have output_scale after fusion"
-    )
 
-    assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
+    assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, \
         "Attention should not have output_block_scale before fusion"
-    )
     if quant_key.dtype == FP8_DTYPE:
-        assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
+        assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, \
             "Attention should not have output_block_scale after FP8 fusion"
-        )
     elif quant_key.dtype == FP4_DTYPE:
-        assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
-            "Attention should have output_block_scale after FP4 fusion"
-        )
+        assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, \
+            "Attention should have output_block_scale after FP4 fusion"  # noqa: E501
 
     # Check that results are close
-    torch.testing.assert_close(result_unfused, result_fused_1, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(result_unfused,
+                               result_fused_1,
+                               atol=1e-2,
+                               rtol=1e-2)

tests/compile/test_noop_elimination.py

@@ -6,12 +6,14 @@ import torch
 
 import vllm
 from vllm.compilation.noop_elimination import NoOpEliminationPass
-from vllm.config import CompilationConfig, CompilationLevel, PassConfig, VllmConfig
+from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,
+                         VllmConfig)
 
 from .backend import TestBackend
 
 
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("dtype",
+                         [torch.float16, torch.bfloat16, torch.float32])
 @pytest.mark.parametrize("num_tokens", [256, 1024])
 @pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size):
@@ -20,6 +22,7 @@ def test_noop_elimination(dtype, num_tokens, hidden_size):
     torch.manual_seed(1)
 
     class Model(torch.nn.Module):
+
         def forward(self, x):
             # Chain of reshapes
             y = x.reshape(-1, 128, 32)
@@ -29,7 +32,7 @@ def test_noop_elimination(dtype, num_tokens, hidden_size):
             # Final reshape that should remain
             b = a.reshape(-1, 128, 32)
             # No-op slice
-            c = b[0 : b.shape[0]]
+            c = b[0:b.shape[0]]
             # The pass should replace the result of this op with `c`.
             d = torch.slice_scatter(
                 torch.ones_like(c),  # Dummy tensor to be scattered into
@@ -40,12 +43,10 @@ def test_noop_elimination(dtype, num_tokens, hidden_size):
             )
             return d
 
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            pass_config=PassConfig(enable_noop=True),
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        pass_config=PassConfig(enable_noop=True),
+    ))
     with vllm.config.set_current_vllm_config(vllm_config):
         noop_pass = NoOpEliminationPass(vllm_config)
 
@@ -81,18 +82,17 @@ def test_non_noop_slice_preserved():
     x = torch.randn(16, 16)
 
     class SliceModel(torch.nn.Module):
+
         def forward(self, x):
             base = x.clone()
             src = torch.ones(15, 16)
             y = torch.slice_scatter(base, src, dim=0, start=0, end=-1)
             return x[0:-1, :], y
 
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            level=CompilationLevel.PIECEWISE,
-            pass_config=PassConfig(enable_noop=True),
-        )
-    )
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        pass_config=PassConfig(enable_noop=True),
+    ))
     with vllm.config.set_current_vllm_config(vllm_config):
         noop_pass = NoOpEliminationPass(vllm_config)
         backend = TestBackend(noop_pass)

tests/compile/test_pass_manager.py

@@ -28,6 +28,7 @@ def test_bad_callable():
 
 # Pass that inherits from InductorPass
 class ProperPass(InductorPass):
+
     def __call__(self, graph: torch.fx.graph.Graph) -> None:
         pass
 
@@ -38,7 +39,8 @@ class ProperPass(InductorPass):
         ProperPass(),
         # Can also wrap callables in CallableInductorPass for compliance
         CallableInductorPass(simple_callable),
-        CallableInductorPass(simple_callable, InductorPass.hash_source(__file__)),
+        CallableInductorPass(simple_callable,
+                             InductorPass.hash_source(__file__))
     ],
 )
 def test_pass_manager_uuid(callable):
@@ -63,9 +65,8 @@ def test_pass_manager_uuid(callable):
 
     # UUID should be different due to config change
     config2 = copy.deepcopy(config)
-    config2.compilation_config.pass_config.enable_fusion = (
-        not config2.compilation_config.pass_config.enable_fusion
-    )
+    config2.compilation_config.pass_config.enable_fusion = not \
+        config2.compilation_config.pass_config.enable_fusion
     pass_manager3 = PostGradPassManager()
     pass_manager3.configure(config2)
     pass_manager3.add(callable)

tests/compile/test_sequence_parallelism.py

@@ -12,20 +12,14 @@ from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
 from vllm.compilation.sequence_parallelism import SequenceParallelismPass
 from vllm.compilation.vllm_inductor_pass import VllmInductorPass
-from vllm.config import (
-    CompilationConfig,
-    DeviceConfig,
-    ModelConfig,
-    PassConfig,
-    VllmConfig,
-)
+from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
+                         PassConfig, VllmConfig)
 from vllm.distributed import tensor_model_parallel_all_reduce
-from vllm.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-)
+from vllm.distributed.parallel_state import (init_distributed_environment,
+                                             initialize_model_parallel)
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    Fp8LinearOp)
 from vllm.platforms import current_platform
 from vllm.utils import update_environment_variables
 
@@ -42,15 +36,16 @@ prompts = [
 
 
 class TestModel(torch.nn.Module):
-    def __init__(
-        self, hidden_size=16, intermediate_size=32, vllm_config: VllmConfig = None
-    ):
+
+    def __init__(self,
+                 hidden_size=16,
+                 intermediate_size=32,
+                 vllm_config: VllmConfig = None):
         super().__init__()
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
         self.gate_proj = torch.nn.Parameter(
-            torch.empty((intermediate_size, hidden_size))
-        )
+            torch.empty((intermediate_size, hidden_size)))
         self.norm = RMSNorm(intermediate_size, 1e-05)
         # Initialize weights
         torch.nn.init.normal_(self.gate_proj, std=0.02)
@@ -58,18 +53,18 @@ class TestModel(torch.nn.Module):
     def forward(self, hidden_states, residual):
         """
         Forward pass implementing the operations in the FX graph
-
+        
         Args:
             hidden_states: Input tensor
             residual: Residual tensor from previous layer
-
+            
         Returns:
             Tuple containing the output tensor
         """
         # Reshape input
         view = hidden_states.reshape(-1, self.hidden_size)
 
-        # matrix multiplication
+        #matrix multiplication
         permute = self.gate_proj.permute(1, 0)
         mm = torch.mm(view, permute)
 
@@ -87,7 +82,7 @@ class TestModel(torch.nn.Module):
     def ops_in_model_after(self):
         return [
             torch.ops.vllm.reduce_scatter.default,
-            torch.ops.vllm.all_gather.default,
+            torch.ops.vllm.all_gather.default
         ]
 
     def ops_in_model(self):
@@ -95,16 +90,18 @@ class TestModel(torch.nn.Module):
 
 
 class TestQuantModel(torch.nn.Module):
-    def __init__(
-        self, hidden_size=16, intermediate_size=32, vllm_config: VllmConfig = None
-    ):
+
+    def __init__(self,
+                 hidden_size=16,
+                 intermediate_size=32,
+                 vllm_config: VllmConfig = None):
         super().__init__()
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
         self.vllm_config = vllm_config
-        self.gate_proj = torch.nn.Parameter(
-            torch.empty((intermediate_size, hidden_size)), requires_grad=False
-        )
+        self.gate_proj = torch.nn.Parameter(torch.empty(
+            (intermediate_size, hidden_size)),
+                                            requires_grad=False)
         self.norm = RMSNorm(intermediate_size, 1e-05)
         # Initialize weights
         torch.nn.init.normal_(self.gate_proj, std=0.02)
@@ -114,24 +111,25 @@ class TestQuantModel(torch.nn.Module):
         self.scale = torch.rand(1, dtype=torch.float32)
         # Create a weight that is compatible with torch._scaled_mm,
         # which expects a column-major layout.
-        self.w = torch.rand(hidden_size, intermediate_size).to(dtype=FP8_DTYPE).t()
+        self.w = torch.rand(hidden_size,
+                            intermediate_size).to(dtype=FP8_DTYPE).t()
         self.wscale = torch.rand(1, dtype=torch.float32)
 
     def forward(self, hidden_states, residual):
         """
         Forward pass implementing the operations in the FX graph
-
+        
         Args:
             hidden_states: Input tensor
             residual: Residual tensor from previous layer
-
+            
         Returns:
             Tuple containing the output tensor
         """
         # Reshape input
         view = hidden_states.reshape(-1, self.hidden_size)
 
-        # matrix multiplication
+        #matrix multiplication
         permute = self.gate_proj.permute(1, 0)
         mm = torch.mm(view, permute)
 
@@ -142,51 +140,45 @@ class TestQuantModel(torch.nn.Module):
         norm_output, residual_output = self.norm(all_reduce, residual)
 
         # scaled_mm with static input quantization
-        fp8_linear_result = self.fp8_linear.apply(
-            norm_output,
-            self.w,
-            self.wscale,
-            input_scale=self.scale.to(norm_output.device),
-        )
+        fp8_linear_result = self.fp8_linear.apply(norm_output,
+                                                  self.w,
+                                                  self.wscale,
+                                                  input_scale=self.scale.to(
+                                                      norm_output.device))
 
         return fp8_linear_result, residual_output
 
     def ops_in_model_before(self):
-        ops_to_remove = [torch.ops.vllm.all_reduce.default]  # Always removed by SP
+        ops_to_remove = [torch.ops.vllm.all_reduce.default
+                         ]  # Always removed by SP
         # The following are only removed if fusion happens
-        if (
-            self.vllm_config
-            and self.vllm_config.compilation_config.pass_config.enable_fusion
-        ):
-            ops_to_remove.extend(
-                [
-                    torch.ops._C.fused_add_rms_norm.default,
-                    torch.ops._C.static_scaled_fp8_quant.default,
-                ]
-            )
+        if self.vllm_config and self.vllm_config.compilation_config \
+            .pass_config.enable_fusion:
+            ops_to_remove.extend([
+                torch.ops._C.fused_add_rms_norm.default,
+                torch.ops._C.static_scaled_fp8_quant.default,
+            ])
         return ops_to_remove
 
     def ops_in_model_after(self):
         ops_to_add = [
             torch.ops.vllm.reduce_scatter.default,
-            torch.ops.vllm.all_gather.default,
+            torch.ops.vllm.all_gather.default
         ]
         # The following is only added if fusion happens
-        if (
-            self.vllm_config
-            and self.vllm_config.compilation_config.pass_config.enable_fusion
-        ):
-            ops_to_add.append(torch.ops._C.fused_add_rms_norm_static_fp8_quant.default)
+        if self.vllm_config and self.vllm_config.compilation_config \
+            .pass_config.enable_fusion:
+            ops_to_add.append(
+                torch.ops._C.fused_add_rms_norm_static_fp8_quant.default)
         return ops_to_add
 
     def ops_in_model(self):
-        if (
-            self.vllm_config
-            and self.vllm_config.compilation_config.pass_config.enable_fusion
-        ):
+        if self.vllm_config and self.vllm_config.compilation_config \
+            .pass_config.enable_fusion:
             # If fusion happens, the fused op is the one
             # we check for (de)functionalization
-            return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
+            return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default
+                    ]  # noqa: E501
         else:
             # If no fusion, the original ops are checked
             return [
@@ -203,47 +195,30 @@ class TestQuantModel(torch.nn.Module):
 @pytest.mark.parametrize("hidden_size", [16])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("enable_fusion", [True, False])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
-def test_sequence_parallelism_pass(
-    test_model_cls: type[torch.nn.Module],
-    batch_size: int,
-    seq_len: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-    enable_fusion: bool,
-):
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
+                    reason="Only test on CUDA")
+def test_sequence_parallelism_pass(test_model_cls: type[torch.nn.Module],
+                                   batch_size: int, seq_len: int,
+                                   hidden_size: int, dtype: torch.dtype,
+                                   enable_fusion: bool):
     num_processes = 2
 
     def run_torch_spawn(fn, nprocs):
         # need to use torch.mp.spawn otherwise will have problems with
         # torch.distributed and cuda
-        torch.multiprocessing.spawn(
-            fn,
-            args=(
-                num_processes,
-                test_model_cls,
-                batch_size,
-                seq_len,
-                hidden_size,
-                dtype,
-                enable_fusion,
-            ),
-            nprocs=nprocs,
-        )
+        torch.multiprocessing.spawn(fn,
+                                    args=(num_processes, test_model_cls,
+                                          batch_size, seq_len, hidden_size,
+                                          dtype, enable_fusion),
+                                    nprocs=nprocs)
 
     run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
 
 
 def sequence_parallelism_pass_on_test_model(
-    local_rank: int,
-    world_size: int,
-    test_model_cls: type[torch.nn.Module],
-    batch_size: int,
-    seq_len: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-    enable_fusion: bool,
-):
+        local_rank: int, world_size: int,
+        test_model_cls: type[torch.nn.Module], batch_size: int, seq_len: int,
+        hidden_size: int, dtype: torch.dtype, enable_fusion: bool):
     current_platform.seed_everything(0)
 
     device = torch.device(f"cuda:{local_rank}")
@@ -251,15 +226,13 @@ def sequence_parallelism_pass_on_test_model(
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "12345",
-        }
-    )
+    update_environment_variables({
+        'RANK': str(local_rank),
+        'LOCAL_RANK': str(local_rank),
+        'WORLD_SIZE': str(world_size),
+        'MASTER_ADDR': 'localhost',
+        'MASTER_PORT': '12345',
+    })
 
     # initialize distributed
     init_distributed_environment()
@@ -267,28 +240,27 @@ def sequence_parallelism_pass_on_test_model(
 
     # configure vllm config for SequenceParallelismPass
     vllm_config = VllmConfig()
-    vllm_config.compilation_config = CompilationConfig(
-        pass_config=PassConfig(
-            enable_sequence_parallelism=True,
-            enable_fusion=enable_fusion,
-            enable_noop=True,
-        )
-    )  # NoOp needed for fusion
+    vllm_config.compilation_config = CompilationConfig(pass_config=PassConfig(
+        enable_sequence_parallelism=True,
+        enable_fusion=enable_fusion,
+        enable_noop=True))  # NoOp needed for fusion
     vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.
     model_name = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP8-e2e"
-    vllm_config.model_config = ModelConfig(
-        model=model_name, trust_remote_code=True, dtype=dtype, seed=42
-    )
+    vllm_config.model_config = ModelConfig(model=model_name,
+                                           trust_remote_code=True,
+                                           dtype=dtype,
+                                           seed=42)
 
     noop_pass = NoOpEliminationPass(vllm_config)
     sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
     func_pass = FixFunctionalizationPass(vllm_config)
     cleanup_pass = PostCleanupPass(vllm_config)
 
-    passes_for_backend: list[VllmInductorPass] = [noop_pass, sequence_parallelism_pass]
+    passes_for_backend: list[VllmInductorPass] = \
+        [noop_pass, sequence_parallelism_pass]
 
     if enable_fusion:
         fusion_pass = RMSNormQuantFusionPass(vllm_config)
@@ -299,9 +271,12 @@ def sequence_parallelism_pass_on_test_model(
     backend_no_func = TestBackend(*passes_for_backend)
     backend_func = TestBackend(*passes_for_backend, func_pass)
 
-    model = test_model_cls(hidden_size, hidden_size * 2, vllm_config=vllm_config)
+    model = test_model_cls(hidden_size,
+                           hidden_size * 2,
+                           vllm_config=vllm_config)
 
-    hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
+    hidden_states = torch.randn((batch_size * seq_len, hidden_size),
+                                dtype=dtype)
     residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
 
     compiled_model_no_func = torch.compile(model, backend=backend_no_func)
@@ -322,7 +297,8 @@ def sequence_parallelism_pass_on_test_model(
     # check if the functionalization pass is applied
     for op in model.ops_in_model():
         find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
-        assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None
+        assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes,
+                                  op) is None  # noqa: E501
 
     # make sure the ops were all de-functionalized
     found = dict()

tests/compile/test_silu_mul_quant_fusion.py

@@ -8,25 +8,20 @@ import torch
 import vllm.envs as envs
 from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
+# yapf conflicts with isort for this block
+# yapf: disable
 from vllm.compilation.activation_quant_fusion import (
-    FUSED_OPS,
-    SILU_MUL_OP,
-    ActivationQuantFusionPass,
-)
+    FUSED_OPS, SILU_MUL_OP, ActivationQuantFusionPass)
+# yapf: enable
 from vllm.compilation.fusion import QUANT_OPS
 from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.compilation.post_cleanup import PostCleanupPass
 from vllm.config import CompilationConfig, PassConfig, VllmConfig
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape,
-    kFp8StaticTensorSym,
-    kNvfp4Quant,
-)
+    GroupShape, kFp8StaticTensorSym, kNvfp4Quant)
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
-    Fp8LinearOp,
-    cutlass_fp8_supported,
-)
+    Fp8LinearOp, cutlass_fp8_supported)
 from vllm.platforms import current_platform
 
 from ..utils import override_cutlass_fp8_supported
@@ -41,6 +36,7 @@ def is_nvfp4_supported():
 
 
 class TestSiluMulFp8QuantModel(torch.nn.Module):
+
     def __init__(self, hidden_size: int, cuda_force_torch: bool, **kwargs):
         super().__init__()
         self.silu_and_mul = SiluAndMul()
@@ -57,7 +53,10 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
 
     def forward(self, x):
         y = self.silu_and_mul(x)
-        x2 = self.fp8_linear.apply(y, self.w, self.wscale, input_scale=self.wscale)
+        x2 = self.fp8_linear.apply(y,
+                                   self.w,
+                                   self.wscale,
+                                   input_scale=self.wscale)
         return x2
 
     def ops_in_model_before(self):
@@ -68,12 +67,11 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
 
 
 class TestSiluMulNvfp4QuantModel(torch.nn.Module):
+
     def __init__(self, hidden_size: int, x: torch.Tensor, **kwargs):
         super().__init__()
         from vllm.compilation.activation_quant_fusion import (
-            silu_and_mul_nvfp4_quant_supported,
-        )
-
+            silu_and_mul_nvfp4_quant_supported)
         assert silu_and_mul_nvfp4_quant_supported
 
         self.silu_and_mul = SiluAndMul()
@@ -90,14 +88,12 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
     def forward(self, x):
         y = self.silu_and_mul(x)
         y_quant, y_block_scale = scaled_fp4_quant(y, self.y_global_scale)
-        out = cutlass_scaled_fp4_mm(
-            a=y_quant,
-            b=self.w,
-            block_scale_a=y_block_scale,
-            block_scale_b=self.w_block_scale,
-            alpha=self.alpha,
-            out_dtype=y.dtype,
-        )
+        out = cutlass_scaled_fp4_mm(a=y_quant,
+                                    b=self.w,
+                                    block_scale_a=y_block_scale,
+                                    block_scale_b=self.w_block_scale,
+                                    alpha=self.alpha,
+                                    out_dtype=y.dtype)
         return out
 
     def ops_in_model_before(self):
@@ -112,24 +108,16 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
 @pytest.mark.parametrize(
     "model_class",
-    cast(
-        list[type],
-        [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
-        if is_nvfp4_supported()
-        else [TestSiluMulFp8QuantModel],
-    ),
-)
+    cast(list[type], [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
+         if is_nvfp4_supported() else [TestSiluMulFp8QuantModel]))
 # cuda_force_torch used to test torch code path on platforms that
 # cutlass_fp8_supported() == True.
-@pytest.mark.parametrize(
-    "cuda_force_torch", [True, False] if cutlass_fp8_supported() else [True]
-)
-@pytest.mark.skipif(
-    envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
-)
-def test_fusion_silu_and_mul_quant(
-    num_tokens, hidden_size, dtype, model_class, cuda_force_torch
-):
+@pytest.mark.parametrize("cuda_force_torch",
+                         [True, False] if cutlass_fp8_supported() else [True])
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
+                    reason="Only test on CUDA and ROCm")
+def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, dtype, model_class,
+                                   cuda_force_torch):
     if model_class == TestSiluMulNvfp4QuantModel and cuda_force_torch:
         pytest.skip("Duplicate tests for NVFP4")
 
@@ -141,13 +129,17 @@ def test_fusion_silu_and_mul_quant(
     # Reshape pass is needed for the fusion pass to work
     config = VllmConfig()
     config.compilation_config = CompilationConfig(
-        pass_config=PassConfig(enable_fusion=True, enable_noop=True)
-    )
+        pass_config=PassConfig(enable_fusion=True, enable_noop=True))
     fusion_pass = ActivationQuantFusionPass(config)
 
-    passes = [NoOpEliminationPass(config), fusion_pass, PostCleanupPass(config)]
+    passes = [
+        NoOpEliminationPass(config), fusion_pass,
+        PostCleanupPass(config)
+    ]
     backend = TestBackend(*passes)
-    model = model_class(hidden_size=hidden_size, cuda_force_torch=cuda_force_torch, x=x)
+    model = model_class(hidden_size=hidden_size,
+                        cuda_force_torch=cuda_force_torch,
+                        x=x)
 
     # First dimension dynamic
     torch._dynamo.mark_dynamic(x, 0)
@@ -163,9 +155,10 @@ def test_fusion_silu_and_mul_quant(
     elif model_class == TestSiluMulNvfp4QuantModel:
         atol, rtol = 1e-1, 1e-1
 
-    torch.testing.assert_close(
-        result[0].to(dtype=dtype), result2[0].to(dtype=dtype), atol=atol, rtol=rtol
-    )
+    torch.testing.assert_close(result[0].to(dtype=dtype),
+                               result2[0].to(dtype=dtype),
+                               atol=atol,
+                               rtol=rtol)
 
     assert fusion_pass.matched_count == 1
 

tests/compile/test_wrapper.py

@@ -10,6 +10,7 @@ from vllm.config import CompilationLevel
 
 
 class MyMod(torch.nn.Module):
+
     def forward(self, x: torch.Tensor, cache: Optional[torch.Tensor] = None):
         if cache is not None:
             return x + cache
@@ -17,12 +18,12 @@ class MyMod(torch.nn.Module):
 
 
 class MyWrapper(TorchCompileWrapperWithCustomDispatcher):
+
     def __init__(self, model):
         self.model = model
         compiled_callable = torch.compile(self.forward, backend="eager")
-        super().__init__(
-            compiled_callable, compilation_level=CompilationLevel.DYNAMO_ONCE
-        )
+        super().__init__(compiled_callable,
+                         compilation_level=CompilationLevel.DYNAMO_ONCE)
 
     def forward(self, x: torch.Tensor, cache: Optional[torch.Tensor] = None):
         # this is the function to be compiled
@@ -53,8 +54,10 @@ def test_torch_compile_wrapper():
 
         # for new input, dispatch to the compiled code directly
         new_x = torch.tensor([3])
-        assert wrapper(new_x, None).item() == 6  # dispatch to the first compiled code
-        assert wrapper(new_x, cache).item() == 5  # dispatch to the second compiled code
+        assert wrapper(new_x,
+                       None).item() == 6  # dispatch to the first compiled code
+        assert wrapper(
+            new_x, cache).item() == 5  # dispatch to the second compiled code
 
     for wrapper in wrappers:
         # make sure they have independent compiled codes

tests/config/test_config_generation.py

@@ -14,9 +14,8 @@ def test_cuda_empty_vs_unset_configs(monkeypatch: pytest.MonkeyPatch):
     """
 
     def create_config():
-        engine_args = EngineArgs(
-            model="deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True
-        )
+        engine_args = EngineArgs(model="deepseek-ai/DeepSeek-V2-Lite",
+                                 trust_remote_code=True)
         return engine_args.create_engine_config()
 
     # Create config with CUDA_VISIBLE_DEVICES set normally
@@ -35,18 +34,16 @@ def test_cuda_empty_vs_unset_configs(monkeypatch: pytest.MonkeyPatch):
     empty_config_dict.pop("instance_id", None)
 
     assert deep_compare(normal_config_dict, empty_config_dict), (
-        'Configs with normal CUDA_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES=""'
-        " should be equivalent"
-    )
+        "Configs with normal CUDA_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES=\"\""
+        " should be equivalent")
 
 
 def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
     # In testing, this method needs to be nested inside as ray does not
     # see the test module.
     def create_config():
-        engine_args = EngineArgs(
-            model="deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True
-        )
+        engine_args = EngineArgs(model="deepseek-ai/DeepSeek-V2-Lite",
+                                 trust_remote_code=True)
         return engine_args.create_engine_config()
 
     config = create_config()
@@ -54,7 +51,6 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
     assert parallel_config.ray_runtime_env is None
 
     import ray
-
     ray.init()
 
     runtime_env = {
@@ -63,13 +59,13 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
         },
     }
 
-    config_ref = ray.remote(create_config).options(runtime_env=runtime_env).remote()
+    config_ref = ray.remote(create_config).options(
+        runtime_env=runtime_env).remote()
 
     config = ray.get(config_ref)
     parallel_config = config.parallel_config
     assert parallel_config.ray_runtime_env is not None
-    assert (
-        parallel_config.ray_runtime_env.env_vars().get("TEST_ENV_VAR") == "test_value"
-    )
+    assert parallel_config.ray_runtime_env.env_vars().get(
+        "TEST_ENV_VAR") == "test_value"
 
     ray.shutdown()

tests/config/test_mp_reducer.py

@@ -16,13 +16,13 @@ def test_mp_reducer(monkeypatch):
     """
 
     # Use V1 AsyncLLM which calls maybe_register_config_serialize_by_value
-    monkeypatch.setenv("VLLM_USE_V1", "1")
+    monkeypatch.setenv('VLLM_USE_V1', '1')
 
     # Ensure transformers_modules is not in sys.modules
-    if "transformers_modules" in sys.modules:
-        del sys.modules["transformers_modules"]
+    if 'transformers_modules' in sys.modules:
+        del sys.modules['transformers_modules']
 
-    with patch("multiprocessing.reducer.register") as mock_register:
+    with patch('multiprocessing.reducer.register') as mock_register:
         engine_args = AsyncEngineArgs(
             model="facebook/opt-125m",
             max_model_len=32,
@@ -36,8 +36,7 @@ def test_mp_reducer(monkeypatch):
         )
 
         assert mock_register.called, (
-            "multiprocessing.reducer.register should have been called"
-        )
+            "multiprocessing.reducer.register should have been called")
 
         vllm_config_registered = False
         for call_args in mock_register.call_args_list:
@@ -46,7 +45,8 @@ def test_mp_reducer(monkeypatch):
                 vllm_config_registered = True
 
                 reducer_func = call_args[0][1]
-                assert callable(reducer_func), "Reducer function should be callable"
+                assert callable(
+                    reducer_func), "Reducer function should be callable"
                 break
 
         assert vllm_config_registered, (

tests/conftest.py

@@ -30,27 +30,23 @@ import torch.nn as nn
 import torch.nn.functional as F
 from huggingface_hub import snapshot_download
 from PIL import Image
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BatchEncoding,
-    BatchFeature,
-)
+from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
+                          BatchEncoding, BatchFeature)
 from transformers.models.auto.auto_factory import _BaseAutoModelClass
 
-from tests.models.utils import TokensTextLogprobs, TokensTextLogprobsPromptLogprobs
+from tests.models.utils import (TokensTextLogprobs,
+                                TokensTextLogprobsPromptLogprobs)
 from vllm import LLM, SamplingParams
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
-from vllm.config.model import ConvertOption, RunnerOption, _get_and_verify_dtype
+from vllm.config.model import (ConvertOption, RunnerOption,
+                               _get_and_verify_dtype)
 from vllm.connections import global_http_connection
-from vllm.distributed import (
-    cleanup_dist_env_and_memory,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
+from vllm.distributed import (cleanup_dist_env_and_memory,
+                              init_distributed_environment,
+                              initialize_model_parallel)
+from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
 from vllm.multimodal.utils import fetch_image
@@ -87,13 +83,12 @@ class ImageAssetPrompts(TypedDict):
 
 
 class ImageTestAssets(list[ImageAsset]):
+
     def __init__(self) -> None:
-        super().__init__(
-            [
-                ImageAsset("stop_sign"),
-                ImageAsset("cherry_blossom"),
-            ]
-        )
+        super().__init__([
+            ImageAsset("stop_sign"),
+            ImageAsset("cherry_blossom"),
+        ])
 
     def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
         """
@@ -110,12 +105,11 @@ class VideoAssetPrompts(TypedDict):
 
 
 class VideoTestAssets(list[VideoAsset]):
+
     def __init__(self) -> None:
-        super().__init__(
-            [
-                VideoAsset("baby_reading"),
-            ]
-        )
+        super().__init__([
+            VideoAsset("baby_reading"),
+        ])
 
     def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
         return [prompts["baby_reading"]]
@@ -127,13 +121,12 @@ class AudioAssetPrompts(TypedDict):
 
 
 class AudioTestAssets(list[AudioAsset]):
+
     def __init__(self) -> None:
-        super().__init__(
-            [
-                AudioAsset("mary_had_lamb"),
-                AudioAsset("winning_call"),
-            ]
-        )
+        super().__init__([
+            AudioAsset("mary_had_lamb"),
+            AudioAsset("winning_call"),
+        ])
 
     def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
         return [prompts["mary_had_lamb"], prompts["winning_call"]]
@@ -228,7 +221,6 @@ def example_system_message() -> str:
 
 class DecoderPromptType(Enum):
     """For encoder/decoder models only."""
-
     CUSTOM = 1
     NONE = 2
     EMPTY_STR = 3
@@ -262,13 +254,15 @@ _R = TypeVar("_R")
 
 
 class HfRunner:
+
     def get_default_device(self):
         from vllm.platforms import current_platform
 
-        return "cpu" if current_platform.is_cpu() else current_platform.device_type
+        return ("cpu"
+                if current_platform.is_cpu() else current_platform.device_type)
 
     def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
-        if x is None or isinstance(x, (bool,)):
+        if x is None or isinstance(x, (bool, )):
             return x
 
         if device is None:
@@ -296,11 +290,8 @@ class HfRunner:
         # Set this to avoid hanging issue
         default_torch_num_threads: Optional[int] = None,
     ) -> None:
-        init_ctx = (
-            nullcontext()
-            if default_torch_num_threads is None
-            else set_default_torch_num_threads(default_torch_num_threads)
-        )
+        init_ctx = (nullcontext() if default_torch_num_threads is None else
+                    set_default_torch_num_threads(default_torch_num_threads))
 
         with init_ctx:
             self._init(
@@ -372,15 +363,14 @@ class HfRunner:
             )
 
             # in case some unquantized custom models are not in same dtype
-            if getattr(model, "quantization_method", None) is None and any(
-                p.dtype != self.dtype for p in model.parameters()
-            ):
+            if (getattr(model, "quantization_method", None) is None
+                    and any(p.dtype != self.dtype
+                            for p in model.parameters())):
                 model = model.to(dtype=self.dtype)
 
-            if (
-                getattr(model, "quantization_method", None) != "bitsandbytes"
-                and len({p.device for p in model.parameters()}) < 2
-            ):
+            if (getattr(model, "quantization_method", None) != "bitsandbytes"
+                    and len({p.device
+                             for p in model.parameters()}) < 2):
                 model = model.to(device=self.device)
 
             self.model = model
@@ -395,7 +385,6 @@ class HfRunner:
         # don't put this import at the top level
         # it will call torch.cuda.device_count()
         from transformers import AutoProcessor  # noqa: F401
-
         self.processor = AutoProcessor.from_pretrained(
             model_name,
             torch_dtype=torch_dtype,
@@ -483,9 +472,10 @@ class HfRunner:
         audios: Optional[PromptAudioInput] = None,
         **kwargs: Any,
     ) -> list[tuple[list[list[int]], list[str]]]:
-        all_inputs = self.get_inputs(
-            prompts, images=images, videos=videos, audios=audios
-        )
+        all_inputs = self.get_inputs(prompts,
+                                     images=images,
+                                     videos=videos,
+                                     audios=audios)
 
         outputs: list[tuple[list[list[int]], list[str]]] = []
         for inputs in all_inputs:
@@ -512,17 +502,16 @@ class HfRunner:
         audios: Optional[PromptAudioInput] = None,
         **kwargs: Any,
     ) -> list[tuple[list[int], str]]:
-        outputs = self.generate(
-            prompts,
-            do_sample=False,
-            max_new_tokens=max_tokens,
-            images=images,
-            videos=videos,
-            audios=audios,
-            **kwargs,
-        )
+        outputs = self.generate(prompts,
+                                do_sample=False,
+                                max_new_tokens=max_tokens,
+                                images=images,
+                                videos=videos,
+                                audios=audios,
+                                **kwargs)
 
-        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
+        return [(output_ids[0], output_str[0])
+                for output_ids, output_str in outputs]
 
     def generate_beam_search(
         self,
@@ -533,22 +522,21 @@ class HfRunner:
         videos: Optional[PromptVideoInput] = None,
         audios: Optional[PromptAudioInput] = None,
     ) -> list[tuple[list[list[int]], list[str]]]:
-        outputs = self.generate(
-            prompts,
-            do_sample=False,
-            max_new_tokens=max_tokens,
-            num_beams=beam_width,
-            num_return_sequences=beam_width,
-            images=images,
-            videos=videos,
-            audios=audios,
-        )
+        outputs = self.generate(prompts,
+                                do_sample=False,
+                                max_new_tokens=max_tokens,
+                                num_beams=beam_width,
+                                num_return_sequences=beam_width,
+                                images=images,
+                                videos=videos,
+                                audios=audios)
 
         for i in range(len(outputs)):
             output_ids, output_str = outputs[i]
             for j in range(len(output_ids)):
                 output_ids[j] = [
-                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
+                    x for x in output_ids[j]
+                    if x != self.tokenizer.pad_token_id
                 ]
             outputs[i] = (output_ids, output_str)
         return outputs
@@ -562,9 +550,10 @@ class HfRunner:
         audios: Optional[PromptAudioInput] = None,
         **kwargs: Any,
     ) -> list[list[torch.Tensor]]:
-        all_inputs = self.get_inputs(
-            prompts, images=images, videos=videos, audios=audios
-        )
+        all_inputs = self.get_inputs(prompts,
+                                     images=images,
+                                     videos=videos,
+                                     audios=audios)
 
         all_logprobs: list[list[torch.Tensor]] = []
         for inputs in all_inputs:
@@ -577,7 +566,8 @@ class HfRunner:
                 return_dict_in_generate=True,
                 **kwargs,
             )
-            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
+            seq_logprobs = self._hidden_states_to_seq_logprobs(
+                output.hidden_states)
             all_logprobs.append(seq_logprobs)
         return all_logprobs
 
@@ -641,9 +631,10 @@ class HfRunner:
         videos: Optional[PromptVideoInput] = None,
         **kwargs: Any,
     ) -> list[TokensTextLogprobs]:
-        all_inputs = self.get_inputs(
-            prompts, images=images, videos=videos, audios=audios
-        )
+        all_inputs = self.get_inputs(prompts,
+                                     images=images,
+                                     videos=videos,
+                                     audios=audios)
 
         all_logprobs: list[list[dict[int, float]]] = []
         all_output_ids: list[list[int]] = []
@@ -663,7 +654,8 @@ class HfRunner:
             (
                 seq_logprobs_lst,
                 output_len,
-            ) = self._hidden_states_to_logprobs(output.hidden_states, num_logprobs)
+            ) = self._hidden_states_to_logprobs(output.hidden_states,
+                                                num_logprobs)
 
             all_logprobs.append(seq_logprobs_lst)
             seq_ids = output.sequences[0]
@@ -673,16 +665,19 @@ class HfRunner:
             all_output_strs.append(self.tokenizer.decode(output_ids))
 
         outputs = zip(all_output_ids, all_output_strs, all_logprobs)
-        return [
-            (output_ids, output_str, output_logprobs)
-            for output_ids, output_str, output_logprobs in outputs
-        ]
+        return [(output_ids, output_str, output_logprobs)
+                for output_ids, output_str, output_logprobs in outputs]
 
-    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
+    def encode(self, prompts: list[str], *args,
+               **kwargs) -> list[list[torch.Tensor]]:
         return self.model.encode(prompts, *args, **kwargs)
 
-    def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor:
-        return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs)
+    def predict(self, prompts: list[list[str]], *args,
+                **kwargs) -> torch.Tensor:
+        return self.model.predict(prompts,
+                                  *args,
+                                  convert_to_tensor=True,
+                                  **kwargs)
 
     def __enter__(self):
         return self
@@ -733,17 +728,10 @@ class VllmRunner:
         default_torch_num_threads: Optional[int] = None,
         **kwargs,
     ) -> None:
-        init_ctx = (
-            nullcontext()
-            if default_torch_num_threads is None
-            else set_default_torch_num_threads(default_torch_num_threads)
-        )
+        init_ctx = (nullcontext() if default_torch_num_threads is None else
+                    set_default_torch_num_threads(default_torch_num_threads))
 
         if not kwargs.get("compilation_config", None):
-            # Note(@tdoublep): This is set to 4 because some tests (e.g., hybrid
-            # model tests) may set max_num_seqs=4. If min cudagraph_capture_size is
-            # set to larger than max_num_seqs, then it will lead to *no* graphs
-            # being captured which can trigger edge cases that we don't handle yet.
             kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
 
         with init_ctx:
@@ -772,25 +760,17 @@ class VllmRunner:
         images: Optional[PromptImageInput] = None,
         videos: Optional[PromptVideoInput] = None,
         audios: Optional[PromptAudioInput] = None,
-    ) -> list[dict[str, Any]]:
-        if any(
-            x is not None and len(x) != len(prompts) for x in [images, videos, audios]
-        ):
+    ) -> list[TextPrompt]:
+
+        if any(x is not None and len(x) != len(prompts)
+               for x in [images, videos, audios]):
             raise ValueError(
-                "All non-None multimodal inputs must have the same length as prompts"
-            )
+                "All non-None multimodal inputs must have the same length as "
+                "prompts")
 
-        inputs = list[dict[str, Any]]()
+        inputs = []
         for i, prompt in enumerate(prompts):
-            prompt_dict = dict[str, Any]()
-            if isinstance(prompt, str):
-                prompt_dict["prompt"] = prompt
-            elif isinstance(prompt, list):
-                prompt_dict["prompt_token_ids"] = prompt
-            else:
-                prompt_dict["prompt_embeds"] = prompt
-
-            multi_modal_data = dict[str, Any]()
+            multi_modal_data = {}
             if images is not None and (image := images[i]) is not None:
                 multi_modal_data["image"] = image
             if videos is not None and (video := videos[i]) is not None:
@@ -798,10 +778,17 @@ class VllmRunner:
             if audios is not None and (audio := audios[i]) is not None:
                 multi_modal_data["audio"] = audio
 
-            if multi_modal_data:
-                prompt_dict["multi_modal_data"] = multi_modal_data
+            text_prompt_kwargs: dict[str, Any] = {
+                "multi_modal_data": multi_modal_data or None
+            }
+            if isinstance(prompt, str):
+                text_prompt_kwargs["prompt"] = prompt
+            elif isinstance(prompt, list):
+                text_prompt_kwargs["prompt_token_ids"] = prompt
+            else:
+                text_prompt_kwargs["prompt_embeds"] = prompt
 
-            inputs.append(prompt_dict)
+            inputs.append(TextPrompt(**text_prompt_kwargs))
 
         return inputs
 
@@ -814,11 +801,14 @@ class VllmRunner:
         audios: Optional[PromptAudioInput] = None,
         **kwargs: Any,
     ) -> list[tuple[list[list[int]], list[str]]]:
-        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
+        inputs = self.get_inputs(prompts,
+                                 images=images,
+                                 videos=videos,
+                                 audios=audios)
 
-        req_outputs = self.llm.generate(
-            inputs, sampling_params=sampling_params, **kwargs
-        )
+        req_outputs = self.llm.generate(inputs,
+                                        sampling_params=sampling_params,
+                                        **kwargs)
 
         outputs: list[tuple[list[list[int]], list[str]]] = []
         for req_output in req_outputs:
@@ -845,9 +835,8 @@ class VllmRunner:
                 output_str = sample.text
                 output_ids = list(sample.token_ids)
                 output_logprobs = sample.logprobs
-            outputs.append(
-                (output_ids, output_str, output_logprobs, req_output.prompt_logprobs)
-            )
+            outputs.append((output_ids, output_str, output_logprobs,
+                            req_output.prompt_logprobs))
         return outputs
 
     def generate_w_logprobs(
@@ -858,22 +847,23 @@ class VllmRunner:
         audios: Optional[PromptAudioInput] = None,
         videos: Optional[PromptVideoInput] = None,
         **kwargs: Any,
-    ) -> Union[list[TokensTextLogprobs], list[TokensTextLogprobsPromptLogprobs]]:
-        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
+    ) -> Union[list[TokensTextLogprobs],
+               list[TokensTextLogprobsPromptLogprobs]]:
+        inputs = self.get_inputs(prompts,
+                                 images=images,
+                                 videos=videos,
+                                 audios=audios)
 
-        req_outputs = self.llm.generate(
-            inputs, sampling_params=sampling_params, **kwargs
-        )
+        req_outputs = self.llm.generate(inputs,
+                                        sampling_params=sampling_params,
+                                        **kwargs)
 
-        toks_str_logsprobs_prompt_logprobs = self._final_steps_generate_w_logprobs(
-            req_outputs
-        )
+        toks_str_logsprobs_prompt_logprobs = (
+            self._final_steps_generate_w_logprobs(req_outputs))
         # Omit prompt logprobs if not required by sampling params
-        return (
-            [x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
-            if sampling_params.prompt_logprobs is None
-            else toks_str_logsprobs_prompt_logprobs
-        )
+        return ([x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
+                if sampling_params.prompt_logprobs is None else
+                toks_str_logsprobs_prompt_logprobs)
 
     def generate_greedy(
         self,
@@ -885,15 +875,14 @@ class VllmRunner:
         **kwargs: Any,
     ) -> list[tuple[list[int], str]]:
         greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-        outputs = self.generate(
-            prompts,
-            greedy_params,
-            images=images,
-            videos=videos,
-            audios=audios,
-            **kwargs,
-        )
-        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
+        outputs = self.generate(prompts,
+                                greedy_params,
+                                images=images,
+                                videos=videos,
+                                audios=audios,
+                                **kwargs)
+        return [(output_ids[0], output_str[0])
+                for output_ids, output_str in outputs]
 
     def generate_greedy_logprobs(
         self,
@@ -907,24 +896,22 @@ class VllmRunner:
         stop_token_ids: Optional[list[int]] = None,
         stop: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> Union[list[TokensTextLogprobs], list[TokensTextLogprobsPromptLogprobs]]:
+    ) -> Union[list[TokensTextLogprobs],
+               list[TokensTextLogprobsPromptLogprobs]]:
         greedy_logprobs_params = SamplingParams(
             temperature=0.0,
             max_tokens=max_tokens,
             logprobs=num_logprobs,
             prompt_logprobs=num_prompt_logprobs,
             stop_token_ids=stop_token_ids,
-            stop=stop,
-        )
+            stop=stop)
 
-        return self.generate_w_logprobs(
-            prompts,
-            greedy_logprobs_params,
-            images=images,
-            audios=audios,
-            videos=videos,
-            **kwargs,
-        )
+        return self.generate_w_logprobs(prompts,
+                                        greedy_logprobs_params,
+                                        images=images,
+                                        audios=audios,
+                                        videos=videos,
+                                        **kwargs)
 
     def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:
         """
@@ -933,9 +920,10 @@ class VllmRunner:
         :param prompts: list of prompts to score
         :return: perplexity score of each prompt
         """
-        outputs = self.generate_greedy_logprobs(
-            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
-        )
+        outputs = self.generate_greedy_logprobs(prompts,
+                                                max_tokens=1,
+                                                num_logprobs=None,
+                                                num_prompt_logprobs=0)
 
         perplexities = []
         for output in outputs:
@@ -964,13 +952,15 @@ class VllmRunner:
         audios: Optional[PromptAudioInput] = None,
         concurrency_limit: Optional[int] = None,
     ) -> list[tuple[list[list[int]], list[str]]]:
-        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
+        inputs = self.get_inputs(prompts,
+                                 images=images,
+                                 videos=videos,
+                                 audios=audios)
 
-        outputs = self.llm.beam_search(
-            inputs,
-            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens),
-            concurrency_limit=concurrency_limit,
-        )
+        outputs = self.llm.beam_search(inputs,
+                                       BeamSearchParams(beam_width=beam_width,
+                                                        max_tokens=max_tokens),
+                                       concurrency_limit=concurrency_limit)
         returned_outputs = []
         for output in outputs:
             token_ids = [x.tokens for x in output.sequences]
@@ -982,16 +972,17 @@ class VllmRunner:
         req_outputs = self.llm.classify(prompts)
         return [req_output.outputs.probs for req_output in req_outputs]
 
-    def embed(
-        self,
-        prompts: list[str],
-        images: Optional[PromptImageInput] = None,
-        videos: Optional[PromptVideoInput] = None,
-        audios: Optional[PromptAudioInput] = None,
-        *args,
-        **kwargs,
-    ) -> list[list[float]]:
-        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
+    def embed(self,
+              prompts: list[str],
+              images: Optional[PromptImageInput] = None,
+              videos: Optional[PromptVideoInput] = None,
+              audios: Optional[PromptAudioInput] = None,
+              *args,
+              **kwargs) -> list[list[float]]:
+        inputs = self.get_inputs(prompts,
+                                 images=images,
+                                 videos=videos,
+                                 audios=audios)
 
         req_outputs = self.llm.embed(inputs, *args, **kwargs)
         return [req_output.outputs.embedding for req_output in req_outputs]
@@ -1036,7 +1027,6 @@ def vllm_runner():
 @pytest.fixture()
 def temporary_enable_log_propagate():
     import logging
-
     logger = logging.getLogger("vllm")
     logger.propagate = True
     yield
@@ -1056,7 +1046,6 @@ def num_gpus_available():
     in current process."""
 
     from vllm.platforms import current_platform
-
     return current_platform.device_count()
 
 
@@ -1070,11 +1059,12 @@ _dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
 def dummy_opt_path():
     json_path = os.path.join(_dummy_opt_path, "config.json")
     if not os.path.exists(_dummy_opt_path):
-        snapshot_download(
-            repo_id="facebook/opt-125m",
-            local_dir=_dummy_opt_path,
-            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
-        )
+        snapshot_download(repo_id="facebook/opt-125m",
+                          local_dir=_dummy_opt_path,
+                          ignore_patterns=[
+                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
+                              "*.msgpack"
+                          ])
         assert os.path.exists(json_path)
         with open(json_path) as f:
             config = json.load(f)
@@ -1088,18 +1078,12 @@ def dummy_opt_path():
 def dummy_llava_path():
     json_path = os.path.join(_dummy_llava_path, "config.json")
     if not os.path.exists(_dummy_llava_path):
-        snapshot_download(
-            repo_id="llava-hf/llava-1.5-7b-hf",
-            local_dir=_dummy_llava_path,
-            ignore_patterns=[
-                "*.bin",
-                "*.bin.index.json",
-                "*.pt",
-                "*.h5",
-                "*.msgpack",
-                "*.safetensors",
-            ],
-        )
+        snapshot_download(repo_id="llava-hf/llava-1.5-7b-hf",
+                          local_dir=_dummy_llava_path,
+                          ignore_patterns=[
+                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
+                              "*.msgpack", "*.safetensors"
+                          ])
         assert os.path.exists(json_path)
         with open(json_path) as f:
             config = json.load(f)
@@ -1113,18 +1097,12 @@ def dummy_llava_path():
 def dummy_gemma2_embedding_path():
     json_path = os.path.join(_dummy_gemma2_embedding_path, "config.json")
     if not os.path.exists(_dummy_gemma2_embedding_path):
-        snapshot_download(
-            repo_id="BAAI/bge-multilingual-gemma2",
-            local_dir=_dummy_gemma2_embedding_path,
-            ignore_patterns=[
-                "*.bin",
-                "*.bin.index.json",
-                "*.pt",
-                "*.h5",
-                "*.msgpack",
-                "*.safetensors",
-            ],
-        )
+        snapshot_download(repo_id="BAAI/bge-multilingual-gemma2",
+                          local_dir=_dummy_gemma2_embedding_path,
+                          ignore_patterns=[
+                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
+                              "*.msgpack", "*.safetensors"
+                          ])
         assert os.path.exists(json_path)
         with open(json_path) as f:
             config = json.load(f)
@@ -1137,9 +1115,10 @@ def dummy_gemma2_embedding_path():
 # Add the flag `--optional` to allow run tests
 # that are marked with @pytest.mark.optional
 def pytest_addoption(parser):
-    parser.addoption(
-        "--optional", action="store_true", default=False, help="run optional test"
-    )
+    parser.addoption("--optional",
+                     action="store_true",
+                     default=False,
+                     help="run optional test")
 
 
 def pytest_collection_modifyitems(config, items):
@@ -1207,6 +1186,7 @@ def _find_free_port() -> int:
 
 
 class LocalAssetServer:
+
     address: str
     port: int
     server: Optional[http.server.ThreadingHTTPServer]
@@ -1221,9 +1201,9 @@ class LocalAssetServer:
     def __enter__(self):
         self.port = _find_free_port()
         self.server = http.server.ThreadingHTTPServer(
-            (self.address, self.port), AssetHandler
-        )
-        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
+            (self.address, self.port), AssetHandler)
+        self.thread = threading.Thread(target=self.server.serve_forever,
+                                       daemon=True)
         self.thread.start()
         return self
 
@@ -1257,7 +1237,7 @@ class LocalAssetServer:
 @pytest.fixture(scope="session")
 def local_asset_server() -> Generator[LocalAssetServer, None, None]:
     """
-    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
+    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port. 
     The server currently servers images at:
     http://127.0.0.1:<port>/<name>.<ext>
     """

tests/cuda/test_cuda_context.py

@@ -13,7 +13,7 @@ from vllm.platforms import current_platform
 def check_cuda_context():
     """Check CUDA driver context status"""
     try:
-        cuda = ctypes.CDLL("libcuda.so")
+        cuda = ctypes.CDLL('libcuda.so')
         device = ctypes.c_int()
         result = cuda.cuCtxGetDevice(ctypes.byref(device))
         return (True, device.value) if result == 0 else (False, None)
@@ -27,11 +27,9 @@ def run_cuda_test_in_thread(device_input, expected_device_id):
         # New thread should have no CUDA context initially
         valid_before, device_before = check_cuda_context()
         if valid_before:
-            return (
-                False,
-                "CUDA context should not exist in new thread, "
-                f"got device {device_before}",
-            )
+            return False, \
+                "CUDA context should not exist in new thread, " \
+                f"got device {device_before}"
 
         # Test setting CUDA context
         current_platform.set_device(device_input)
@@ -41,7 +39,8 @@ def run_cuda_test_in_thread(device_input, expected_device_id):
         if not valid_after:
             return False, "CUDA context should be valid after set_cuda_context"
         if device_id != expected_device_id:
-            return False, f"Expected device {expected_device_id}, got {device_id}"
+            return False, \
+                f"Expected device {expected_device_id}, got {device_id}"
 
         return True, "Success"
     except Exception as e:
@@ -51,30 +50,30 @@ def run_cuda_test_in_thread(device_input, expected_device_id):
 class TestSetCudaContext:
     """Test suite for the set_cuda_context function."""
 
-    @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
-    @pytest.mark.parametrize(
-        argnames="device_input,expected_device_id",
-        argvalues=[
-            (0, 0),
-            (torch.device("cuda:0"), 0),
-            ("cuda:0", 0),
-        ],
-        ids=["int", "torch_device", "string"],
-    )
-    def test_set_cuda_context_parametrized(self, device_input, expected_device_id):
+    @pytest.mark.skipif(not current_platform.is_cuda(),
+                        reason="CUDA not available")
+    @pytest.mark.parametrize(argnames="device_input,expected_device_id",
+                             argvalues=[
+                                 (0, 0),
+                                 (torch.device('cuda:0'), 0),
+                                 ('cuda:0', 0),
+                             ],
+                             ids=["int", "torch_device", "string"])
+    def test_set_cuda_context_parametrized(self, device_input,
+                                           expected_device_id):
         """Test setting CUDA context in isolated threads."""
         with ThreadPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(
-                run_cuda_test_in_thread, device_input, expected_device_id
-            )
+            future = executor.submit(run_cuda_test_in_thread, device_input,
+                                     expected_device_id)
             success, message = future.result(timeout=30)
         assert success, message
 
-    @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
+    @pytest.mark.skipif(not current_platform.is_cuda(),
+                        reason="CUDA not available")
     def test_set_cuda_context_invalid_device_type(self):
         """Test error handling for invalid device type."""
         with pytest.raises(ValueError, match="Expected a cuda device"):
-            current_platform.set_device(torch.device("cpu"))
+            current_platform.set_device(torch.device('cpu'))
 
 
 if __name__ == "__main__":

tests/detokenizer/test_disable_detokenization.py

@@ -17,16 +17,20 @@ def test_computed_prefix_blocks(model: str):
     prompt = (
         "You are a helpful assistant. How do I build a car from cardboard and "
         "paper clips? Is there an easy to follow video tutorial available "
-        "online for free?"
-    )
+        "online for free?")
 
     llm = LLM(model=model)
-    sampling_params = SamplingParams(max_tokens=10, temperature=0.0, detokenize=False)
+    sampling_params = SamplingParams(max_tokens=10,
+                                     temperature=0.0,
+                                     detokenize=False)
 
-    outputs_no_detokenization = llm.generate(prompt, sampling_params)[0].outputs[0]
+    outputs_no_detokenization = llm.generate(prompt,
+                                             sampling_params)[0].outputs[0]
     sampling_params.detokenize = True
-    outputs_with_detokenization = llm.generate(prompt, sampling_params)[0].outputs[0]
+    outputs_with_detokenization = llm.generate(prompt,
+                                               sampling_params)[0].outputs[0]
 
-    assert outputs_no_detokenization.text == ""
-    assert outputs_with_detokenization.text != ""
-    assert outputs_no_detokenization.token_ids == outputs_with_detokenization.token_ids
+    assert outputs_no_detokenization.text == ''
+    assert outputs_with_detokenization.text != ''
+    assert outputs_no_detokenization.token_ids == \
+        outputs_with_detokenization.token_ids

tests/detokenizer/test_min_tokens.py

@@ -8,17 +8,15 @@ from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.detokenizer import FastIncrementalDetokenizer
 
-PROMPT = "Hello, my name is Lee, and I'm a student in the " + "college of engineering"
+PROMPT = "Hello, my name is Lee, and I'm a student in the " + \
+         "college of engineering"
 
 
-@pytest.mark.parametrize(
-    "min_tokens,stop,truth",
-    [
-        (0, None, " is Lee, and I'm a student in the college of engineering"),
-        (0, "e", " is L"),
-        (5, "e", " is Lee, and I'm a stud"),
-    ],
-)
+@pytest.mark.parametrize("min_tokens,stop,truth", [
+    (0, None, " is Lee, and I'm a student in the college of engineering"),
+    (0, "e", " is L"),
+    (5, "e", " is Lee, and I'm a stud"),
+])
 def test_min_tokens_with_stop(min_tokens: int, stop: str, truth: str):
     """Test for a specific min_tokens and stop.
 
@@ -33,18 +31,16 @@ def test_min_tokens_with_stop(min_tokens: int, stop: str, truth: str):
         stop=stop,
         min_tokens=min_tokens,
     )
-    request = EngineCoreRequest(
-        request_id="",
-        prompt_token_ids=prompt_token_ids,
-        mm_features=None,
-        sampling_params=params,
-        pooling_params=None,
-        eos_token_id=None,
-        arrival_time=0.0,
-        lora_request=None,
-        cache_salt=None,
-        data_parallel_rank=None,
-    )
+    request = EngineCoreRequest(request_id="",
+                                prompt_token_ids=prompt_token_ids,
+                                mm_features=None,
+                                sampling_params=params,
+                                pooling_params=None,
+                                eos_token_id=None,
+                                arrival_time=0.0,
+                                lora_request=None,
+                                cache_salt=None,
+                                data_parallel_rank=None)
 
     detokenizer = FastIncrementalDetokenizer(tokenizer, request)
 

tests/detokenizer/test_stop_reason.py

@@ -31,39 +31,34 @@ def test_stop_reason(vllm_model, example_prompts):
     llm = vllm_model.llm
 
     # test stop token
-    outputs = llm.generate(
-        example_prompts,
-        sampling_params=SamplingParams(
-            ignore_eos=True,
-            seed=SEED,
-            max_tokens=MAX_TOKENS,
-            stop_token_ids=[stop_token_id],
-        ),
-    )
+    outputs = llm.generate(example_prompts,
+                           sampling_params=SamplingParams(
+                               ignore_eos=True,
+                               seed=SEED,
+                               max_tokens=MAX_TOKENS,
+                               stop_token_ids=[stop_token_id]))
     for output in outputs:
         output = output.outputs[0]
         assert output.finish_reason == "stop"
         assert output.stop_reason == stop_token_id
 
     # test stop string
-    outputs = llm.generate(
-        example_prompts,
-        sampling_params=SamplingParams(
-            ignore_eos=True, seed=SEED, max_tokens=MAX_TOKENS, stop="."
-        ),
-    )
+    outputs = llm.generate(example_prompts,
+                           sampling_params=SamplingParams(
+                               ignore_eos=True,
+                               seed=SEED,
+                               max_tokens=MAX_TOKENS,
+                               stop="."))
     for output in outputs:
         output = output.outputs[0]
         assert output.finish_reason == "stop"
         assert output.stop_reason == STOP_STR
 
     # test EOS token
-    outputs = llm.generate(
-        example_prompts,
-        sampling_params=SamplingParams(seed=SEED, max_tokens=MAX_TOKENS),
-    )
+    outputs = llm.generate(example_prompts,
+                           sampling_params=SamplingParams(
+                               seed=SEED, max_tokens=MAX_TOKENS))
     for output in outputs:
         output = output.outputs[0]
         assert output.finish_reason == "length" or (
-            output.finish_reason == "stop" and output.stop_reason is None
-        )
+            output.finish_reason == "stop" and output.stop_reason is None)

tests/detokenizer/test_stop_string_while_stop_model_terminates.py

@@ -14,6 +14,7 @@ def include_stop_str_in_output(request):
 
 
 class _DummyDetokenizer(BaseIncrementalDetokenizer):
+
     def __init__(self, request: EngineCoreRequest):
         super().__init__(request)
 
@@ -26,8 +27,7 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0):
     params = SamplingParams(
         stop=stop,
         include_stop_str_in_output=include_stop_str_in_output,
-        min_tokens=min_tokens,
-    )
+        min_tokens=min_tokens)
     # Keep other fields minimal for unit test purposes.
     req = EngineCoreRequest(
         request_id="test",
@@ -44,25 +44,26 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0):
     return req
 
 
-def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: bool):
+def test_stop_string_while_stop_token_terminates(
+        include_stop_str_in_output: bool):
     """
     This test verifies that the detokenizer correctly handles the case where
     the generated token sequence contains both:
     - a stop token
     - an <eos> token
-
+    
     The detokenizer should respect the stop string and truncate the output
     accordingly.
-
+    
     Imagine the following sequence:
     - "abcdeZ" is generated, where "Z" is the <eos> token.
     - "cd" is the stop string.
-
+    
     If include_stop_str_in_output=False, the detokenizer should truncate the
     output to "ab" because the stop string "cd" is excluded.
     If include_stop_str_in_output=True, the detokenizer should include the stop
     string "cd" in the output, resulting in "abcd".
-
+    
 
     This verifies the behavioral change introduced in BaseIncrementalDetokenizer
     where stop-string evaluation occurs before the early-return on
@@ -77,9 +78,8 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo
     token_ids = [ord(c) for c in generated_text]
 
     # Create a request with the stop string and initialize the detokenizer.
-    req = _make_request(
-        stop=[stop_string], include_stop_str_in_output=include_stop_str_in_output
-    )
+    req = _make_request(stop=[stop_string],
+                        include_stop_str_in_output=include_stop_str_in_output)
     detok = _DummyDetokenizer(req)
 
     # Simulate that the last token ('Z') is a stop token (stop_terminated=True).
@@ -99,4 +99,5 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo
 
     # get_next_output_text should return the full text when finished=True.
     # (Buffering only applies during streaming when finished=False.)
-    assert detok.get_next_output_text(finished=True, delta=False) == expected_text
+    assert detok.get_next_output_text(finished=True,
+                                      delta=False) == expected_text

tests/detokenizer/test_stop_strings.py

@@ -11,14 +11,12 @@ MODEL = "meta-llama/llama-2-7b-hf"
 MAX_TOKENS = 200
 
 
-def _test_stopping(
-    llm: LLM,
-    expected_output: str,
-    expected_reason: Any,
-    stop: Optional[list[str]] = None,
-    stop_token_ids: Optional[list[int]] = None,
-    include_in_output: bool = False,
-) -> None:
+def _test_stopping(llm: LLM,
+                   expected_output: str,
+                   expected_reason: Any,
+                   stop: Optional[list[str]] = None,
+                   stop_token_ids: Optional[list[int]] = None,
+                   include_in_output: bool = False) -> None:
     output = llm.generate(
         "A story about vLLM:\n",
         SamplingParams(
@@ -27,8 +25,7 @@ def _test_stopping(
             stop=stop,
             stop_token_ids=stop_token_ids,
             include_stop_str_in_output=include_in_output,
-        ),
-    )[0].outputs[0]
+        ))[0].outputs[0]
 
     assert output is not None
     assert output.text == expected_output
@@ -36,21 +33,17 @@ def _test_stopping(
 
 
 def _stop_basic(llm):
-    _test_stopping(
-        llm,
-        stop=["."],
-        include_in_output=False,
-        expected_output="VLLM is a 100% volunteer organization",
-        expected_reason=".",
-    )
+    _test_stopping(llm,
+                   stop=["."],
+                   include_in_output=False,
+                   expected_output="VLLM is a 100% volunteer organization",
+                   expected_reason=".")
 
-    _test_stopping(
-        llm,
-        stop=["."],
-        include_in_output=True,
-        expected_output="VLLM is a 100% volunteer organization.",
-        expected_reason=".",
-    )
+    _test_stopping(llm,
+                   stop=["."],
+                   include_in_output=True,
+                   expected_output="VLLM is a 100% volunteer organization.",
+                   expected_reason=".")
 
 
 def _stop_multi_tokens(llm):
@@ -59,54 +52,45 @@ def _stop_multi_tokens(llm):
         stop=["group of peo", "short"],
         include_in_output=False,
         expected_output="VLLM is a 100% volunteer organization. We are a ",
-        expected_reason="group of peo",
-    )
+        expected_reason="group of peo")
 
     _test_stopping(
         llm,
         stop=["group of peo", "short"],
         include_in_output=True,
-        expected_output="VLLM is a 100% volunteer organization. We are a group of peo",
-        expected_reason="group of peo",
-    )
+        expected_output=
+        "VLLM is a 100% volunteer organization. We are a group of peo",
+        expected_reason="group of peo")
 
 
 def _stop_partial_token(llm):
-    _test_stopping(
-        llm,
-        stop=["gani"],
-        include_in_output=False,
-        expected_output="VLLM is a 100% volunteer or",
-        expected_reason="gani",
-    )
+    _test_stopping(llm,
+                   stop=["gani"],
+                   include_in_output=False,
+                   expected_output="VLLM is a 100% volunteer or",
+                   expected_reason="gani")
 
-    _test_stopping(
-        llm,
-        stop=["gani"],
-        include_in_output=True,
-        expected_output="VLLM is a 100% volunteer organi",
-        expected_reason="gani",
-    )
+    _test_stopping(llm,
+                   stop=["gani"],
+                   include_in_output=True,
+                   expected_output="VLLM is a 100% volunteer organi",
+                   expected_reason="gani")
 
 
 def _stop_token_id(llm):
     # token id 13013 => " organization"
 
-    _test_stopping(
-        llm,
-        stop_token_ids=[13013],
-        include_in_output=False,
-        expected_output="VLLM is a 100% volunteer",
-        expected_reason=13013,
-    )
+    _test_stopping(llm,
+                   stop_token_ids=[13013],
+                   include_in_output=False,
+                   expected_output="VLLM is a 100% volunteer",
+                   expected_reason=13013)
 
-    _test_stopping(
-        llm,
-        stop_token_ids=[13013],
-        include_in_output=True,
-        expected_output="VLLM is a 100% volunteer organization",
-        expected_reason=13013,
-    )
+    _test_stopping(llm,
+                   stop_token_ids=[13013],
+                   include_in_output=True,
+                   expected_output="VLLM is a 100% volunteer organization",
+                   expected_reason=13013)
 
 
 @pytest.mark.skip_global_cleanup

tests/distributed/conftest.py

@@ -111,7 +111,8 @@ class MockSubscriber:
         self.last_seq = -1
         self.decoder = msgspec.msgpack.Decoder(type=decode_type)
 
-    def receive_one(self, timeout=1000) -> Union[tuple[int, SampleBatch], None]:
+    def receive_one(self,
+                    timeout=1000) -> Union[tuple[int, SampleBatch], None]:
         """Receive a single message with timeout"""
         if not self.sub.poll(timeout):
             return None
@@ -134,7 +135,8 @@ class MockSubscriber:
 
         self.replay_sockets[socket_idx].send(start_seq.to_bytes(8, "big"))
 
-    def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
+    def receive_replay(self,
+                       socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
         """Receive replayed messages from a specific replay socket"""
         if not self.replay_sockets:
             raise ValueError("Replay sockets not initialized")

tests/distributed/test_ca_buffer_sharing.py

@@ -12,8 +12,7 @@ import torch.distributed as dist
 
 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from vllm.distributed.device_communicators.custom_all_reduce import (  # noqa
-    CustomAllreduce,
-)
+    CustomAllreduce)
 
 # create a cpu process group for communicating metadata (ipc handle)
 dist.init_process_group(backend="gloo")
@@ -53,8 +52,7 @@ for p in pointers:
         assert ord(host_data[i]) == byte_value, (
             f"Rank {rank} failed"
             f" to verify buffer {p}. Expected {byte_value}, "
-            f"got {ord(host_data[i])}"
-        )
+            f"got {ord(host_data[i])}")
 
 print(f"Rank {rank} verified all buffers")
 

tests/distributed/test_comm_ops.py

@@ -13,19 +13,13 @@ import pytest
 import ray
 import torch
 
-from vllm.distributed import (
-    broadcast_tensor_dict,
-    get_pp_group,
-    tensor_model_parallel_all_gather,
-    tensor_model_parallel_all_reduce,
-    tensor_model_parallel_reduce_scatter,
-)
+from vllm.distributed import (broadcast_tensor_dict, get_pp_group,
+                              tensor_model_parallel_all_gather,
+                              tensor_model_parallel_all_reduce,
+                              tensor_model_parallel_reduce_scatter)
 
-from ..utils import (
-    init_test_distributed_environment,
-    multi_gpu_test,
-    multi_process_parallel,
-)
+from ..utils import (init_test_distributed_environment, multi_gpu_test,
+                     multi_process_parallel)
 
 
 @ray.remote(num_gpus=1, max_calls=1)
@@ -43,11 +37,12 @@ def all_reduce_test_worker(
 
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
     num_elements = 8
     all_tensors = [
-        torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1)
-        for r in range(tp_size)
+        torch.arange(num_elements, dtype=torch.float32, device="cuda") *
+        (r + 1) for r in range(tp_size)
     ]
     expected = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
     t = all_tensors[rank % tp_size]
@@ -56,31 +51,28 @@ def all_reduce_test_worker(
 
 
 @ray.remote(num_gpus=1, max_calls=1)
-def reduce_scatter_test_worker(
-    monkeypatch: pytest.MonkeyPatch,
-    tp_size: int,
-    pp_size: int,
-    rank: int,
-    distributed_init_port: str,
-):
+def reduce_scatter_test_worker(monkeypatch: pytest.MonkeyPatch, tp_size: int,
+                               pp_size: int, rank: int,
+                               distributed_init_port: str):
     # it is important to delete the CUDA_VISIBLE_DEVICES environment variable
     # so that each worker can see all the GPUs
     # they will be able to set the device to the correct GPU
     monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
 
     num_elements = 8
     all_tensors = [
-        torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1)
-        for r in range(tp_size)
+        torch.arange(num_elements, dtype=torch.float32, device="cuda") *
+        (r + 1) for r in range(tp_size)
     ]
 
     index = rank % tp_size
     partition_size = num_elements // tp_size
     all_reduce = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
-    expected = all_reduce[index * partition_size : (index + 1) * partition_size]
+    expected = all_reduce[index * partition_size:(index + 1) * partition_size]
     t = all_tensors[index]
     t = tensor_model_parallel_reduce_scatter(t, 0)
     torch.testing.assert_close(t, expected)
@@ -100,7 +92,8 @@ def all_gather_test_worker(
     monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
     num_dimensions = 3
     tensor_size = list(range(2, num_dimensions + 2))
     total_size = 1
@@ -108,10 +101,8 @@ def all_gather_test_worker(
         total_size *= s
     for all_gather_dimension in range(num_dimensions):
         all_tensors = [
-            torch.arange(total_size, dtype=torch.float32, device="cuda").reshape(
-                tensor_size
-            )
-            * (r + 1)
+            torch.arange(total_size, dtype=torch.float32,
+                         device="cuda").reshape(tensor_size) * (r + 1)
             for r in range(tp_size)
         ]
         expected = torch.cat(all_tensors, dim=all_gather_dimension)
@@ -134,7 +125,8 @@ def broadcast_tensor_dict_test_worker(
     monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
     test_dict = {
         # device tensor
         "a": torch.arange(8, dtype=torch.float32, device="cuda"),
@@ -142,7 +134,10 @@ def broadcast_tensor_dict_test_worker(
         "b": torch.arange(16, dtype=torch.int8, device="cpu"),
         "c": "test",
         "d": [1, 2, 3],
-        "e": {"a": 1, "b": 2},
+        "e": {
+            "a": 1,
+            "b": 2
+        },
         # empty tensor
         "f": torch.tensor([], dtype=torch.float32, device="cuda"),
     }
@@ -171,7 +166,8 @@ def send_recv_tensor_dict_test_worker(
     monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
 
     test_dict = {
         # device tensor
@@ -180,7 +176,10 @@ def send_recv_tensor_dict_test_worker(
         "b": torch.arange(16, dtype=torch.int8, device="cpu"),
         "c": "test",
         "d": [1, 2, 3],
-        "e": {"a": 1, "b": 2},
+        "e": {
+            "a": 1,
+            "b": 2
+        },
         # empty tensor
         "f": torch.tensor([], dtype=torch.float32, device="cuda"),
     }
@@ -212,7 +211,8 @@ def send_recv_test_worker(
     monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+    init_test_distributed_environment(tp_size, pp_size, rank,
+                                      distributed_init_port)
 
     size = 64
     test_tensor = torch.arange(64, dtype=torch.float32, device="cuda")
@@ -229,10 +229,10 @@ def send_recv_test_worker(
 
 @multi_gpu_test(num_gpus=2)
 @pytest.mark.parametrize("tp_size", [2])
-@pytest.mark.parametrize(
-    "test_target",
-    [all_reduce_test_worker, all_gather_test_worker, broadcast_tensor_dict_test_worker],
-)
+@pytest.mark.parametrize("test_target", [
+    all_reduce_test_worker, all_gather_test_worker,
+    broadcast_tensor_dict_test_worker
+])
 def test_multi_process_tensor_parallel(
     monkeypatch: pytest.MonkeyPatch,
     tp_size: int,
@@ -244,8 +244,7 @@ def test_multi_process_tensor_parallel(
 @multi_gpu_test(num_gpus=2)
 @pytest.mark.parametrize("pp_size", [2])
 @pytest.mark.parametrize(
-    "test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker]
-)
+    "test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker])
 def test_multi_process_pipeline_parallel(
     monkeypatch: pytest.MonkeyPatch,
     pp_size: int,
@@ -257,16 +256,11 @@ def test_multi_process_pipeline_parallel(
 @multi_gpu_test(num_gpus=4)
 @pytest.mark.parametrize("tp_size", [2])
 @pytest.mark.parametrize("pp_size", [2])
-@pytest.mark.parametrize(
-    "test_target",
-    [
-        send_recv_test_worker,
-        send_recv_tensor_dict_test_worker,
-        all_reduce_test_worker,
-        all_gather_test_worker,
-        broadcast_tensor_dict_test_worker,
-    ],
-)
+@pytest.mark.parametrize("test_target", [
+    send_recv_test_worker, send_recv_tensor_dict_test_worker,
+    all_reduce_test_worker, all_gather_test_worker,
+    broadcast_tensor_dict_test_worker
+])
 def test_multi_process_tensor_parallel_pipeline_parallel(
     tp_size: int,
     pp_size: int,

tests/distributed/test_context_parallel.py

@@ -7,7 +7,6 @@ WARNING: This test runs in both single-node (4 GPUs) and multi-node
  all workers in a node other than the head node, which can cause the test
  to fail.
 """
-
 import json
 import os
 from dataclasses import dataclass
@@ -57,8 +56,7 @@ class CPTestSettings:
             raise ValueError(
                 f"Length mismatch: distributed_backends "
                 f"({len(self.distributed_backends)}) != "
-                f"vllm_major_versions ({len(self.vllm_major_versions)})"
-            )
+                f"vllm_major_versions ({len(self.vllm_major_versions)})")
 
     @staticmethod
     def detailed(
@@ -76,39 +74,29 @@ class CPTestSettings:
                 for dcp_multiplier in [0.5, 1]:
                     for chunked_prefill_val in [True]:
                         parallel_setups.append(
-                            ParallelSetup(
-                                tp_size=tp_base,
-                                pp_size=pp_multiplier * pp_base,
-                                dcp_size=int(dcp_multiplier * tp_base),
-                                eager_mode=eager_mode_val,
-                                chunked_prefill=chunked_prefill_val,
-                            )
-                        )
+                            ParallelSetup(tp_size=tp_base,
+                                          pp_size=pp_multiplier * pp_base,
+                                          dcp_size=int(dcp_multiplier *
+                                                       tp_base),
+                                          eager_mode=eager_mode_val,
+                                          chunked_prefill=chunked_prefill_val))
         return CPTestSettings(
             parallel_setups=parallel_setups,
             distributed_backends=["mp"],
             vllm_major_versions=["1"],
             runner=runner,
-            test_options=CPTestOptions(
-                multi_node_only=multi_node_only, load_format=load_format
-            ),
+            test_options=CPTestOptions(multi_node_only=multi_node_only,
+                                       load_format=load_format),
         )
 
     def iter_params(self, model_id: str):
         opts = self.test_options
 
         for parallel_setup in self.parallel_setups:
-            for backend, vllm_major_version in zip(
-                self.distributed_backends, self.vllm_major_versions
-            ):
-                yield (
-                    model_id,
-                    parallel_setup,
-                    backend,
-                    vllm_major_version,
-                    self.runner,
-                    opts,
-                )
+            for backend, vllm_major_version in zip(self.distributed_backends,
+                                                   self.vllm_major_versions):
+                yield (model_id, parallel_setup, backend, vllm_major_version,
+                       self.runner, opts)
 
 
 def _compare_cp_with_tp(
@@ -160,10 +148,8 @@ def _compare_cp_with_tp(
     if num_gpus_available < tp_size * pp_size:
         pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
     if VLLM_MULTI_NODE and distributed_backend == "mp":
-        pytest.skip(
-            "Skipping multi-node pipeline parallel test for "
-            "multiprocessing distributed backend"
-        )
+        pytest.skip("Skipping multi-node pipeline parallel test for "
+                    "multiprocessing distributed backend")
     if multi_node_only and not VLLM_MULTI_NODE:
         pytest.skip("Not in multi-node setting")
 
@@ -192,7 +178,8 @@ def _compare_cp_with_tp(
         common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
 
     cp_env = tp_env = {
-        "VLLM_USE_V1": vllm_major_version,  # Note(hc): DCP only support V1 engine only
+        "VLLM_USE_V1":
+        vllm_major_version,  # Note(hc): DCP only support V1 engine only
     }
 
     cp_args = [
@@ -218,15 +205,13 @@ def _compare_cp_with_tp(
     ]
 
     try:
-        compare_two_settings(
-            model_id,
-            cp_args,
-            tp_args,
-            cp_env,
-            tp_env,
-            method=method,
-            max_wait_seconds=720,
-        )
+        compare_two_settings(model_id,
+                             cp_args,
+                             tp_args,
+                             cp_env,
+                             tp_env,
+                             method=method,
+                             max_wait_seconds=720)
     except Exception:
         testing_ray_compiled_graph = cp_env is not None
         if testing_ray_compiled_graph and vllm_major_version == "0":
@@ -239,10 +224,9 @@ def _compare_cp_with_tp(
 
 CP_TEXT_GENERATION_MODELS = {
     # [MLA attention only]
-    "deepseek-ai/DeepSeek-V2-Lite-Chat": [
-        CPTestSettings.detailed(),
-        CPTestSettings.detailed(tp_base=2),
-    ],
+    "deepseek-ai/DeepSeek-V2-Lite-Chat":
+    [CPTestSettings.detailed(),
+     CPTestSettings.detailed(tp_base=2)],
 }
 
 CP_TEST_MODELS = [
@@ -253,19 +237,11 @@ CP_TEST_MODELS = [
 
 
 @pytest.mark.parametrize(
-    (
-        "model_id",
-        "parallel_setup",
-        "distributed_backend",
-        "vllm_major_version",
-        "runner",
-        "test_options",
-    ),
+    ("model_id", "parallel_setup", "distributed_backend", "vllm_major_version",
+     "runner", "test_options"),
     [
-        params
-        for model_id, settings in CP_TEXT_GENERATION_MODELS.items()
-        for setting in settings
-        for params in setting.iter_params(model_id)
+        params for model_id, settings in CP_TEXT_GENERATION_MODELS.items()
+        for setting in settings for params in setting.iter_params(model_id)
         if model_id in CP_TEST_MODELS
     ],
 )
@@ -279,14 +255,12 @@ def test_cp_generation(
     test_options: CPTestOptions,
     num_gpus_available,
 ):
-    _compare_cp_with_tp(
-        model_id,
-        parallel_setup,
-        distributed_backend,
-        vllm_major_version,
-        runner,
-        test_options,
-        num_gpus_available,
-        method="generate",
-        is_multimodal=False,
-    )
+    _compare_cp_with_tp(model_id,
+                        parallel_setup,
+                        distributed_backend,
+                        vllm_major_version,
+                        runner,
+                        test_options,
+                        num_gpus_available,
+                        method="generate",
+                        is_multimodal=False)

tests/distributed/test_custom_all_reduce.py

@@ -8,14 +8,12 @@ import ray
 import torch
 import torch.distributed as dist
 
-from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
+from vllm.distributed.communication_op import (  # noqa
+    tensor_model_parallel_all_reduce)
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
 
-from ..utils import (
-    ensure_model_parallel_initialized,
-    init_test_distributed_environment,
-    multi_process_parallel,
-)
+from ..utils import (ensure_model_parallel_initialized,
+                     init_test_distributed_environment, multi_process_parallel)
 
 random.seed(42)
 test_sizes = [random.randint(1024, 2048 * 1024) for _ in range(8)]
@@ -35,7 +33,8 @@ def graph_allreduce(
         m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
         device = torch.device(f"cuda:{rank}")
         torch.cuda.set_device(device)
-        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+        init_test_distributed_environment(tp_size, pp_size, rank,
+                                          distributed_init_port)
         ensure_model_parallel_initialized(tp_size, pp_size)
         group = get_tp_group().device_group
 
@@ -61,15 +60,18 @@ def graph_allreduce(
             for dtype in [torch.float32, torch.float16, torch.bfloat16]:
                 with graph_capture(device=device) as graph_capture_context:
                     # use integers so result matches NCCL exactly
-                    inp1 = torch.randint(
-                        1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
-                    )
-                    inp2 = torch.randint(
-                        1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
-                    )
+                    inp1 = torch.randint(1,
+                                         16, (sz, ),
+                                         dtype=dtype,
+                                         device=torch.cuda.current_device())
+                    inp2 = torch.randint(1,
+                                         16, (sz, ),
+                                         dtype=dtype,
+                                         device=torch.cuda.current_device())
                     torch.cuda.synchronize()
                     graph = torch.cuda.CUDAGraph()
-                    with torch.cuda.graph(graph, stream=graph_capture_context.stream):
+                    with torch.cuda.graph(graph,
+                                          stream=graph_capture_context.stream):
                         for i in range(num_communication):
                             out1 = tensor_model_parallel_all_reduce(inp1)
                             # the input buffer is immediately modified to test
@@ -94,7 +96,8 @@ def eager_allreduce(
         m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
         device = torch.device(f"cuda:{rank}")
         torch.cuda.set_device(device)
-        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+        init_test_distributed_environment(tp_size, pp_size, rank,
+                                          distributed_init_port)
 
         # we use the first group to communicate once
         # and the second group to communicate twice
@@ -129,4 +132,5 @@ def test_custom_allreduce(
     world_size = tp_size * pipeline_parallel_size
     if world_size > torch.cuda.device_count():
         pytest.skip("Not enough GPUs to run the test.")
-    multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
+    multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size,
+                           test_target)

tests/distributed/test_distributed_oot.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from ..entrypoints.openai.test_oot_registration import run_and_test_dummy_opt_api_server
+from ..entrypoints.openai.test_oot_registration import (
+    run_and_test_dummy_opt_api_server)
 
 
 def test_distributed_oot(dummy_opt_path: str):

tests/distributed/test_eplb_algo.py

@@ -10,12 +10,10 @@ from vllm.distributed.eplb.rebalance_algo import rebalance_experts
 def test_basic_rebalance():
     """Test basic rebalancing functionality"""
     # Example from https://github.com/deepseek-ai/eplb
-    weight = torch.tensor(
-        [
-            [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
-            [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
-        ]
-    )
+    weight = torch.tensor([
+        [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
+        [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
+    ])
 
     num_layers = weight.shape[0]
     num_replicas = 16
@@ -23,49 +21,45 @@ def test_basic_rebalance():
     num_nodes = 2
     num_gpus = 8
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify output shapes
     assert phy2log.shape == (
         2,
         16,
     ), f"Expected `phy2log` shape (2, 16), got {phy2log.shape}"
-    assert log2phy.shape[0] == 2, (
-        f"Expected `log2phy` first dimension 2, got {log2phy.shape[0]}"
-    )
-    assert log2phy.shape[1] == 12, (
-        f"Expected `log2phy` second dimension 12, got {log2phy.shape[1]}"
-    )
+    assert (log2phy.shape[0] == 2
+            ), f"Expected `log2phy` first dimension 2, got {log2phy.shape[0]}"
+    assert (
+        log2phy.shape[1] == 12
+    ), f"Expected `log2phy` second dimension 12, got {log2phy.shape[1]}"
     assert logcnt.shape == (
         2,
         12,
     ), f"Expected `logcnt` shape (2, 12), got {logcnt.shape}"
 
     # Verify physical to logical expert mapping range is correct
-    assert torch.all(phy2log >= 0) and torch.all(phy2log < 12), (
-        "Physical to logical mapping should be in range [0, 12)"
-    )
+    assert torch.all(phy2log >= 0) and torch.all(
+        phy2log < 12), "Physical to logical mapping should be in range [0, 12)"
 
     # Verify expert count reasonableness
-    assert torch.all(logcnt >= 1), "Each logical expert should have at least 1 replica"
-    assert torch.sum(logcnt, dim=1).sum() == num_replicas * num_layers, (
-        f"Total replicas should be {num_replicas * num_layers}"
-    )
+    assert torch.all(
+        logcnt >= 1), "Each logical expert should have at least 1 replica"
+    assert (
+        torch.sum(logcnt, dim=1).sum() == num_replicas *
+        num_layers), f"Total replicas should be {num_replicas * num_layers}"
 
     # Verify expected output
-    expected_phy2log = torch.tensor(
-        [
-            [5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
-            [7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1],
-        ]
-    )
+    expected_phy2log = torch.tensor([
+        [5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
+        [7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1],
+    ])
     assert torch.all(phy2log == expected_phy2log)
 
-    expected_logcnt = torch.tensor(
-        [[1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1], [1, 2, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1]]
-    )
+    expected_logcnt = torch.tensor([[1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1],
+                                    [1, 2, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1]])
     assert torch.all(logcnt == expected_logcnt)
 
 
@@ -77,9 +71,9 @@ def test_single_gpu_case():
     num_nodes = 1
     num_gpus = 1
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify shapes
     assert phy2log.shape == (1, 4)
@@ -99,19 +93,19 @@ def test_equal_weights():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify shapes
     assert phy2log.shape == (1, 8)
     assert logcnt.shape == (1, 8)
 
     # With equal weights, each expert should have exactly one replica
-    assert torch.all(logcnt == 1), (
-        "With equal weights and no replication, "
-        "each expert should have exactly 1 replica"
-    )
+    assert torch.all(
+        logcnt == 1
+    ), "With equal weights and no replication, " \
+       "each expert should have exactly 1 replica"
 
 
 def test_extreme_weight_imbalance():
@@ -122,37 +116,35 @@ def test_extreme_weight_imbalance():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify shapes
     assert phy2log.shape == (1, 12)
     assert logcnt.shape == (1, 8)
 
     # Expert with highest weight (index 0) should have more replicas
-    assert logcnt[0, 0] > logcnt[0, 1], (
-        "Expert with highest weight should have more replicas"
-    )
+    assert (
+        logcnt[0, 0]
+        > logcnt[0, 1]), "Expert with highest weight should have more replicas"
 
 
 def test_multiple_layers():
     """Test multiple layers case"""
-    weight = torch.tensor(
-        [
-            [10, 20, 30, 40, 50, 60],  # First layer
-            [60, 50, 40, 30, 20, 10],  # Second layer (opposite weight pattern)
-            [25, 25, 25, 25, 25, 25],  # Third layer (equal weights)
-        ]
-    )
+    weight = torch.tensor([
+        [10, 20, 30, 40, 50, 60],  # First layer
+        [60, 50, 40, 30, 20, 10],  # Second layer (opposite weight pattern)
+        [25, 25, 25, 25, 25, 25],  # Third layer (equal weights)
+    ])
     num_replicas = 8
     num_groups = 2
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify shapes
     assert phy2log.shape == (3, 8)
@@ -160,12 +152,12 @@ def test_multiple_layers():
 
     # Verify expert allocation is reasonable for each layer
     for layer in range(3):
-        assert torch.all(phy2log[layer] >= 0) and torch.all(phy2log[layer] < 6), (
-            f"Layer {layer} physical to logical mappingshould be in range [0, 6)"
-        )
-        assert torch.sum(logcnt[layer]) == num_replicas, (
-            f"Layer {layer} total replicas should be {num_replicas}"
-        )
+        assert torch.all(phy2log[layer] >= 0) and torch.all(
+            phy2log[layer] < 6
+        ), f"Layer {layer} physical to logical mapping" \
+            "should be in range [0, 6)"
+        assert (torch.sum(logcnt[layer]) == num_replicas
+                ), f"Layer {layer} total replicas should be {num_replicas}"
 
 
 def test_parameter_validation():
@@ -187,19 +179,17 @@ def test_parameter_validation():
 
 def test_small_scale_hierarchical():
     """Test small-scale hierarchical load balancing"""
-    weight = torch.tensor(
-        [
-            [100, 50, 200, 75, 150, 25, 300, 80],  # 8 experts
-        ]
-    )
+    weight = torch.tensor([
+        [100, 50, 200, 75, 150, 25, 300, 80],  # 8 experts
+    ])
     num_replicas = 12
     num_groups = 4  # 4 groups, 2 experts each
     num_nodes = 2  # 2 nodes
     num_gpus = 4  # 4 GPUs
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Verify basic constraints
     assert phy2log.shape == (1, 12)
@@ -209,9 +199,8 @@ def test_small_scale_hierarchical():
 
     # Expert with highest weight should have more replicas
     max_weight_expert = torch.argmax(weight[0])
-    assert logcnt[0, max_weight_expert] >= 2, (
-        "Highest weight expert should have multiple replicas"
-    )
+    assert (logcnt[0, max_weight_expert]
+            >= 2), "Highest weight expert should have multiple replicas"
 
 
 def test_global_load_balance_fallback():
@@ -224,9 +213,9 @@ def test_global_load_balance_fallback():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Should work normally, just using global load balancing strategy
     assert phy2log.shape == (1, 8)
@@ -246,9 +235,9 @@ def test_device_compatibility(device):
     num_nodes = 1
     num_gpus = 2
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
 
     # Function will convert to CPU internally, but should handle different
     # device inputs normally
@@ -261,8 +250,7 @@ def test_additional_cases():
 
     # Test case 1: Large-scale distributed setup
     weight1 = torch.tensor(
-        [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
-    )
+        [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]])
     phy2log1, log2phy1, logcnt1 = rebalance_experts(weight1, 24, 8, 4, 8)
 
     assert phy2log1.shape == (1, 24)
@@ -270,12 +258,10 @@ def test_additional_cases():
     assert torch.sum(logcnt1) == 24
 
     # Test case 2: Different weight distributions
-    weight2 = torch.tensor(
-        [
-            [200, 150, 100, 50, 25, 12],  # Decreasing weights
-            [12, 25, 50, 100, 150, 200],  # Increasing weights
-        ]
-    )
+    weight2 = torch.tensor([
+        [200, 150, 100, 50, 25, 12],  # Decreasing weights
+        [12, 25, 50, 100, 150, 200],  # Increasing weights
+    ])
     phy2log2, log2phy2, logcnt2 = rebalance_experts(weight2, 10, 3, 1, 2)
 
     assert phy2log2.shape == (2, 10)
@@ -288,21 +274,19 @@ def test_additional_cases():
 
 
 if __name__ == "__main__":
-    weight = torch.tensor(
-        [
-            [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
-            [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
-        ]
-    )
+    weight = torch.tensor([
+        [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
+        [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
+    ])
 
     num_replicas = 16
     num_groups = 4
     num_nodes = 2
     num_gpus = 8
 
-    phy2log, log2phy, logcnt = rebalance_experts(
-        weight, num_replicas, num_groups, num_nodes, num_gpus
-    )
+    phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
+                                                 num_groups, num_nodes,
+                                                 num_gpus)
     print(phy2log)
 
     test_basic_rebalance()

tests/distributed/test_eplb_execute.py

@@ -9,12 +9,11 @@ import pytest
 import torch
 import torch.distributed
 
-from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
-from vllm.distributed.parallel_state import (
-    ensure_model_parallel_initialized,
-    get_tp_group,
-    init_distributed_environment,
-)
+from vllm.distributed.eplb.rebalance_execute import (
+    rearrange_expert_weights_inplace)
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
+                                             get_tp_group,
+                                             init_distributed_environment)
 from vllm.utils import update_environment_variables
 
 
@@ -23,13 +22,13 @@ def distributed_run(fn, world_size):
     processes: list[multiprocessing.Process] = []
     for i in range(number_of_processes):
         env: dict[str, str] = {}
-        env["RANK"] = str(i)
-        env["LOCAL_RANK"] = str(i)
-        env["WORLD_SIZE"] = str(number_of_processes)
-        env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
-        env["MASTER_ADDR"] = "localhost"
-        env["MASTER_PORT"] = "12345"
-        p = multiprocessing.Process(target=fn, args=(env,))
+        env['RANK'] = str(i)
+        env['LOCAL_RANK'] = str(i)
+        env['WORLD_SIZE'] = str(number_of_processes)
+        env['LOCAL_WORLD_SIZE'] = str(number_of_processes)
+        env['MASTER_ADDR'] = 'localhost'
+        env['MASTER_PORT'] = '12345'
+        p = multiprocessing.Process(target=fn, args=(env, ))
         processes.append(p)
         p.start()
 
@@ -46,7 +45,7 @@ def worker_fn_wrapper(fn):
     # and update the environment variables in the function
     def wrapped_fn(env):
         update_environment_variables(env)
-        local_rank = os.environ["LOCAL_RANK"]
+        local_rank = os.environ['LOCAL_RANK']
         device = torch.device(f"cuda:{local_rank}")
         torch.cuda.set_device(device)
         init_distributed_environment()
@@ -61,20 +60,20 @@ def worker_fn_wrapper(fn):
 
 
 def create_expert_indices_with_redundancy(
-    num_layers: int,
-    num_logical_experts: int,
-    total_physical_experts: int,
-    redundancy_config: list[int],  # redundancy for each logical expert
+        num_layers: int,
+        num_logical_experts: int,
+        total_physical_experts: int,
+        redundancy_config: list[int],  # redundancy for each logical expert
 ) -> torch.Tensor:
     """
     Create expert indices with redundancy.
-
+    
     Args:
         num_layers: number of layers
         num_logical_experts: number of logical experts
         total_physical_experts: total number of physical experts
         redundancy_config: redundancy for each logical expert
-
+    
     Returns:
         indices: Shape (num_layers, total_physical_experts)
     """
@@ -107,11 +106,11 @@ def create_expert_weights(
 ) -> list[list[torch.Tensor]]:
     """
     Create fake expert weights tensor for testing.
-
+    
     Use `arange` to generate predictable weights values, based on logical
     expert ID.
     All replicas of the same logical expert should have the same weights.
-
+    
     Args:
         physical_to_logical_mapping: Shape (num_layers, num_local_experts)
             mapping[layer, physical_pos] = logical_expert_id
@@ -121,27 +120,27 @@ def create_expert_weights(
     for layer in range(num_layers):
         layer_weights = []
         for weight_idx, hidden_size in enumerate(hidden_sizes):
-            weight_tensor = torch.zeros(
-                num_local_experts, hidden_size, device=device, dtype=torch.float32
-            )
+            weight_tensor = torch.zeros(num_local_experts,
+                                        hidden_size,
+                                        device=device,
+                                        dtype=torch.float32)
 
             for local_expert in range(num_local_experts):
                 # Get the logical expert ID for this physical expert
                 global_pos = rank * num_local_experts + local_expert
                 logical_expert_id = physical_to_logical_mapping[
-                    layer, global_pos
-                ].item()
+                    layer, global_pos].item()
 
                 # Generate weights based on logical expert ID
                 # (so that all replicas of the same logical expert have the
                 # same weights)
-                base_value = logical_expert_id * 1000 + layer * 100 + weight_idx * 10
-                weight_tensor[local_expert] = torch.arange(
-                    base_value,
-                    base_value + hidden_size,
-                    device=device,
-                    dtype=torch.float32,
-                )
+                base_value = (logical_expert_id * 1000 + layer * 100 +
+                              weight_idx * 10)
+                weight_tensor[local_expert] = torch.arange(base_value,
+                                                           base_value +
+                                                           hidden_size,
+                                                           device=device,
+                                                           dtype=torch.float32)
 
             layer_weights.append(weight_tensor)
         expert_weights.append(layer_weights)
@@ -183,15 +182,12 @@ def verify_expert_weights_after_shuffle(
 
                 # Check if the weights are correct
                 actual_weights = weight_tensor[local_expert]
-                expected_base = (
-                    expected_logical_expert * 1000 + layer * 100 + weight_idx * 10
-                )
-                expected_weights = torch.arange(
-                    expected_base,
-                    expected_base + hidden_size,
-                    device=actual_weights.device,
-                    dtype=actual_weights.dtype,
-                )
+                expected_base = (expected_logical_expert * 1000 + layer * 100 +
+                                 weight_idx * 10)
+                expected_weights = torch.arange(expected_base,
+                                                expected_base + hidden_size,
+                                                device=actual_weights.device,
+                                                dtype=actual_weights.dtype)
 
                 torch.testing.assert_close(
                     actual_weights,
@@ -199,8 +195,7 @@ def verify_expert_weights_after_shuffle(
                     msg=f"Layer {layer}, weight {weight_idx},"
                     f"local expert {local_expert}: "
                     f"weights do not match. "
-                    f"Expected logical expert {expected_logical_expert}",
-                )
+                    f"Expected logical expert {expected_logical_expert}")
 
 
 def verify_redundant_experts_have_same_weights(
@@ -227,23 +222,23 @@ def verify_redundant_experts_have_same_weights(
                 total_physical_experts,
                 hidden_size,
                 device=expert_weights[layer][weight_idx].device,
-                dtype=expert_weights[layer][weight_idx].dtype,
-            )
+                dtype=expert_weights[layer][weight_idx].dtype)
 
             # Use all_gather to collect expert weights from current node
             # expert_weights[layer][weight_idx] shape:
             # [num_local_experts, hidden_size]
             local_weights = expert_weights[layer][
-                weight_idx
-            ]  # [num_local_experts, hidden_size]
+                weight_idx]  # [num_local_experts, hidden_size]
 
             # Split tensor along dim 0 into a list for all_gather
-            gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0)
+            gathered_weights_list = torch.chunk(gathered_weights,
+                                                world_size,
+                                                dim=0)
 
             torch.distributed.all_gather(
                 # Output list: each element corresponds to one rank's weights
                 list(gathered_weights_list),
-                local_weights,  # Input: current rank's local weights
+                local_weights  # Input: current rank's local weights
             )
 
             all_weights.append(gathered_weights)
@@ -271,8 +266,7 @@ def verify_redundant_experts_have_same_weights(
                         msg=f"Layer {layer}, weight {weight_idx},"
                         f"logical expert {logical_expert_id}: "
                         f"Physical expert {physical_pos} has different weights"
-                        f"than expected",
-                    )
+                        f"than expected")
 
 
 @pytest.mark.parametrize(
@@ -296,11 +290,10 @@ def verify_redundant_experts_have_same_weights(
         # 4 GPU, 8 experts per GPU
         # 16 logical experts, 32 physical experts, 16 redundant experts
         (4, 8, 8, 16),
-    ],
-)
-def test_rearrange_expert_weights_with_redundancy(
-    world_size, num_layers, num_local_experts, num_logical_experts
-):
+    ])
+def test_rearrange_expert_weights_with_redundancy(world_size, num_layers,
+                                                  num_local_experts,
+                                                  num_logical_experts):
     """Test the functionality of rearranging expert weights with redundancy."""
 
     if torch.cuda.device_count() < world_size:
@@ -311,8 +304,8 @@ def test_rearrange_expert_weights_with_redundancy(
         # Initialize model parallel (using tensor parallel as an entrypoint
         # to expert parallel)
         ensure_model_parallel_initialized(
-            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
-        )
+            tensor_model_parallel_size=world_size,
+            pipeline_model_parallel_size=1)
 
         ep_group = get_tp_group().cpu_group
         ep_rank = torch.distributed.get_rank()
@@ -323,9 +316,8 @@ def test_rearrange_expert_weights_with_redundancy(
         hidden_sizes = [32, 64]  # Two different weight matrices
 
         # Create old expert indices (with redundancy)
-        redundancy_config = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
+        redundancy_config = create_redundancy_config(num_logical_experts,
+                                                     total_physical_experts)
 
         old_indices = create_expert_indices_with_redundancy(
             num_layers,
@@ -336,8 +328,7 @@ def test_rearrange_expert_weights_with_redundancy(
 
         # Create new expert indices (with redundancy)
         new_redundancy_config = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
+            num_logical_experts, total_physical_experts)
         new_indices = create_expert_indices_with_redundancy(
             num_layers,
             num_logical_experts,
@@ -346,9 +337,9 @@ def test_rearrange_expert_weights_with_redundancy(
         )
 
         # Create expert weights
-        expert_weights = create_expert_weights(
-            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
-        )
+        expert_weights = create_expert_weights(num_layers, num_local_experts,
+                                               hidden_sizes, ep_rank, device,
+                                               old_indices)
 
         # Execute weight rearrangement
         rearrange_expert_weights_inplace(
@@ -392,8 +383,8 @@ def test_rearrange_expert_weights_no_change(world_size):
     @worker_fn_wrapper
     def worker_fn():
         ensure_model_parallel_initialized(
-            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
-        )
+            tensor_model_parallel_size=world_size,
+            pipeline_model_parallel_size=1)
 
         ep_group = get_tp_group().cpu_group
         ep_rank = torch.distributed.get_rank()
@@ -410,12 +401,12 @@ def test_rearrange_expert_weights_no_change(world_size):
 
         # Same indices - no change
         indices = create_expert_indices_with_redundancy(
-            num_layers, num_logical_experts, total_physical_experts, redundancy_config
-        )
+            num_layers, num_logical_experts, total_physical_experts,
+            redundancy_config)
 
-        expert_weights = create_expert_weights(
-            num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
-        )
+        expert_weights = create_expert_weights(num_layers, num_local_experts,
+                                               hidden_sizes, ep_rank, device,
+                                               indices)
 
         # Save original weights
         original_weights = []
@@ -431,8 +422,7 @@ def test_rearrange_expert_weights_no_change(world_size):
             indices,  # Same indices
             expert_weights,
             ep_group,
-            is_profile=False,
-        )
+            is_profile=False)
 
         # Verify that the weights have not changed
         for layer in range(num_layers):
@@ -440,8 +430,8 @@ def test_rearrange_expert_weights_no_change(world_size):
                 torch.testing.assert_close(
                     expert_weights[layer][weight_idx],
                     original_weights[layer][weight_idx],
-                    msg=f"Layer {layer}, weight {weight_idx} should remain unchanged",
-                )
+                    msg=f"Layer {layer}, weight {weight_idx} should remain "
+                    f"unchanged")
 
     distributed_run(worker_fn, world_size)
 
@@ -456,8 +446,8 @@ def test_rearrange_expert_weights_profile_mode(world_size):
     @worker_fn_wrapper
     def worker_fn():
         ensure_model_parallel_initialized(
-            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
-        )
+            tensor_model_parallel_size=world_size,
+            pipeline_model_parallel_size=1)
 
         ep_group = get_tp_group().cpu_group
         ep_rank = torch.distributed.get_rank()
@@ -470,23 +460,21 @@ def test_rearrange_expert_weights_profile_mode(world_size):
         hidden_sizes = [32]
 
         # Create different index distributions
-        old_redundancy = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
-        new_redundancy = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
+        old_redundancy = create_redundancy_config(num_logical_experts,
+                                                  total_physical_experts)
+        new_redundancy = create_redundancy_config(num_logical_experts,
+                                                  total_physical_experts)
 
         old_indices = create_expert_indices_with_redundancy(
-            num_layers, num_logical_experts, total_physical_experts, old_redundancy
-        )
+            num_layers, num_logical_experts, total_physical_experts,
+            old_redundancy)
         new_indices = create_expert_indices_with_redundancy(
-            num_layers, num_logical_experts, total_physical_experts, new_redundancy
-        )
+            num_layers, num_logical_experts, total_physical_experts,
+            new_redundancy)
 
-        expert_weights = create_expert_weights(
-            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
-        )
+        expert_weights = create_expert_weights(num_layers, num_local_experts,
+                                               hidden_sizes, ep_rank, device,
+                                               old_indices)
 
         # Save original weights
         original_weights = []
@@ -502,7 +490,7 @@ def test_rearrange_expert_weights_profile_mode(world_size):
             new_indices,
             expert_weights,
             ep_group,
-            is_profile=True,  # Profile mode
+            is_profile=True  # Profile mode
         )
 
         # In profile mode, the weights should remain unchanged
@@ -511,7 +499,6 @@ def test_rearrange_expert_weights_profile_mode(world_size):
                 torch.testing.assert_close(
                     expert_weights[layer][weight_idx],
                     original_weights[layer][weight_idx],
-                    msg="In profile mode, the weights should remain unchanged",
-                )
+                    msg="In profile mode, the weights should remain unchanged")
 
     distributed_run(worker_fn, world_size)

tests/distributed/test_events.py

@@ -6,29 +6,24 @@ import time
 import msgspec
 import pytest
 
-from vllm.distributed.kv_events import (
-    EventBatch,
-    EventPublisherFactory,
-    NullEventPublisher,
-)
+from vllm.distributed.kv_events import (EventBatch, EventPublisherFactory,
+                                        NullEventPublisher)
 
 DP_RANK = 0
 
 
 class EventSample(
-    msgspec.Struct,
-    tag=True,  # type: ignore
-    array_like=True,  # type: ignore
+        msgspec.Struct,
+        tag=True,  # type: ignore
+        array_like=True  # type: ignore
 ):
     """Test event for publisher testing"""
-
     id: int
     value: str
 
 
 class SampleBatch(EventBatch):
     """Test event batch for publisher testing"""
-
     events: list[EventSample]
 
 
@@ -49,8 +44,10 @@ def test_basic_publishing(publisher, subscriber):
 
     seq, received = result
     assert seq == 0, "Sequence number mismatch"
-    assert received.ts == pytest.approx(test_batch.ts, abs=0.1), "Timestamp mismatch"
-    assert len(received.events) == len(test_batch.events), "Number of events mismatch"
+    assert received.ts == pytest.approx(test_batch.ts,
+                                        abs=0.1), ("Timestamp mismatch")
+    assert len(received.events) == len(
+        test_batch.events), ("Number of events mismatch")
 
     for i, event in enumerate(received.events):
         assert event.id == i, "Event id mismatch"
@@ -91,9 +88,9 @@ def test_replay_mechanism(publisher, subscriber):
     assert len(replayed) > 0, "No replayed messages received"
     seqs = [seq for seq, _ in replayed]
     assert all(seq >= 10 for seq in seqs), "Replayed messages not in order"
-    assert seqs == list(range(min(seqs), max(seqs) + 1)), (
-        "Replayed messages not consecutive"
-    )
+    assert seqs == list(range(min(seqs),
+                              max(seqs) +
+                              1)), ("Replayed messages not consecutive")
 
 
 def test_buffer_limit(publisher, subscriber, publisher_config):
@@ -129,7 +126,6 @@ def test_topic_filtering(publisher_config):
     pub = EventPublisherFactory.create(publisher_config, DP_RANK)
 
     from .conftest import MockSubscriber
-
     sub_foo = MockSubscriber(publisher_config.endpoint, None, "foo")
     sub_bar = MockSubscriber(publisher_config.endpoint, None, "bar")
 
@@ -141,13 +137,11 @@ def test_topic_filtering(publisher_config):
 
         foo_received = [sub_foo.receive_one(timeout=200) for _ in range(3)]
         assert all(msg is not None for msg in foo_received), (
-            "Subscriber with matching topic should receive messages"
-        )
+            "Subscriber with matching topic should receive messages")
 
         bar_received = [sub_bar.receive_one(timeout=200) for _ in range(3)]
         assert all(msg is None for msg in bar_received), (
-            "Subscriber with non-matching topic should receive no messages"
-        )
+            "Subscriber with non-matching topic should receive no messages")
     finally:
         pub.shutdown()
         sub_foo.close()
@@ -184,7 +178,8 @@ def test_high_volume(publisher, subscriber):
 
     publisher_thread.join()
 
-    assert len(received) >= num_batches * 0.9, "We should have received most messages"
+    assert len(received) >= num_batches * 0.9, (
+        "We should have received most messages")
 
     seqs = [seq for seq, _ in received]
     assert sorted(seqs) == seqs, "Sequence numbers should be in order"
@@ -214,15 +209,13 @@ def test_data_parallel_rank_tagging(publisher_config):
         # For TCP endpoints: tcp://localhost:5557 -> tcp://localhost:5557, tcp://localhost:5558
         expected_endpoint_0 = base_endpoint  # rank 0 gets port + 0 = same port
         expected_endpoint_1 = base_endpoint.replace(
-            ":5557", ":5558"
-        )  # rank 1 gets port + 1
+            ":5557", ":5558")  # rank 1 gets port + 1
     else:
         # For inproc endpoints: inproc://test -> inproc://test_dp0, inproc://test_dp1
         expected_endpoint_0 = base_endpoint  # rank 0 gets base
         expected_endpoint_1 = base_endpoint + "_dp1"  # rank 1 gets _dp1
 
     from .conftest import MockSubscriber
-
     sub_0 = MockSubscriber(expected_endpoint_0, None, publisher_config.topic)
     sub_1 = MockSubscriber(expected_endpoint_1, None, publisher_config.topic)
 
@@ -248,15 +241,15 @@ def test_data_parallel_rank_tagging(publisher_config):
 
         # Verify DP rank tagging
         assert received_0.data_parallel_rank == 0, (
-            f"Expected DP rank 0, got {received_0.data_parallel_rank}"
-        )
+            f"Expected DP rank 0, got {received_0.data_parallel_rank}")
         assert received_1.data_parallel_rank == 1, (
-            f"Expected DP rank 1, got {received_1.data_parallel_rank}"
-        )
+            f"Expected DP rank 1, got {received_1.data_parallel_rank}")
 
         # Verify event content is correct
-        assert len(received_0.events) == 2, "Wrong number of events from rank 0"
-        assert len(received_1.events) == 3, "Wrong number of events from rank 1"
+        assert len(
+            received_0.events) == 2, "Wrong number of events from rank 0"
+        assert len(
+            received_1.events) == 3, "Wrong number of events from rank 1"
 
     finally:
         pub_0.shutdown()