Add pydocstyle to the ruff rules (#52136)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 16:44:29 +00:00
parent 804e5377aa
commit 8c1557a79c
2002 changed files with 15024 additions and 14269 deletions
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Simplified batch specification grammar for attention benchmarks.
+"""Simplified batch specification grammar for attention benchmarks.

 Grammar (underscore-separated segments):
  Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
@@ -24,6 +23,7 @@ Examples:
  q4s1k            -> [(4, 1024)]              # 4-token extend (spec decode)
  2q1k_32q1s1k     -> [(1024, 1024)] * 2 + [(1, 1024)] * 32  # Mixed batch
  16q4s1k          -> [(4, 1024)] * 16         # 16 spec decode requests
+
 """

 from collections import Counter
@@ -71,8 +71,7 @@ def _parse_size(size_str: str, k_suffix: str) -> int:


 def parse_batch_spec(spec: str) -> list[BatchRequest]:
-    """
-    Parse batch specification string into list of BatchRequest objects.
+    """Parse batch specification string into list of BatchRequest objects.

    Grammar: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?

@@ -84,6 +83,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]:

    Raises:
        ValueError: If spec format is invalid
+
    """
    requests = []

@@ -103,8 +103,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]:


 def format_batch_spec(requests: list[BatchRequest]) -> str:
-    """
-    Format list of BatchRequest into human-readable string.
+    """Format list of BatchRequest into human-readable string.

    Groups requests by type and provides counts and sizes.

@@ -113,6 +112,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str:

    Returns:
        Formatted string describing the batch
+
    """
    kinds = {
        "prefill": [],
@@ -157,8 +157,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str:


 def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
-    """
-    Reorder requests for FlashInfer: decode first, then prefill.
+    """Reorder requests for FlashInfer: decode first, then prefill.

    FlashInfer expects decode requests before prefill requests for
    optimal performance.
@@ -168,6 +167,7 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:

    Returns:
        Reordered list with decode requests first
+
    """
    decodes = [r for r in requests if r.is_decode]
    non_decodes = [r for r in requests if not r.is_decode]
@@ -177,14 +177,14 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
 def split_by_type(
    requests: list[BatchRequest],
 ) -> dict[str, list[BatchRequest]]:
-    """
-    Split requests by type for analysis.
+    """Split requests by type for analysis.

    Args:
        requests: List of BatchRequest

    Returns:
        Dict with keys: 'decode', 'prefill', 'extend'
+
    """
    result = {
        "decode": [],
@@ -204,14 +204,14 @@ def split_by_type(


 def get_batch_stats(requests: list[BatchRequest]) -> dict:
-    """
-    Compute statistics about a batch.
+    """Compute statistics about a batch.

    Args:
        requests: List of BatchRequest

    Returns:
        Dict with batch statistics
+
    """
    by_type = split_by_type(requests)

@@ -232,8 +232,7 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:


 def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
-    """
-    Classify a batch spec into a type string.
+    """Classify a batch spec into a type string.

    Args:
        batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
@@ -241,6 +240,7 @@ def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:

    Returns:
        Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
+
    """
    requests = parse_batch_spec(batch_spec)

@@ -2,8 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Universal vLLM Attention Benchmark
+"""Universal vLLM Attention Benchmark.

 Benchmark any attention backend with the extended grammar.
 Supports standard attention (Flash/Triton/FlashInfer) and MLA backends.
@@ -23,6 +22,7 @@ Examples:

    # Parameter sweep (YAML config - recommended)
    python benchmark.py --config configs/cutlass_numsplits.yaml
+
 """

 import argparse
@@ -84,8 +84,7 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:


 def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
-    """
-    Run a single benchmark with proper backend selection.
+    """Run a single benchmark with proper backend selection.

    Args:
        config: BenchmarkConfig with backend, batch_spec, and model params
@@ -93,6 +92,7 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:

    Returns:
        BenchmarkResult (may have error field set on failure)
+
    """
    try:
        if is_mla_backend(config.backend):
@@ -119,8 +119,7 @@ def run_model_parameter_sweep(
    sweep: ModelParameterSweep,
    console: Console,
 ) -> list[BenchmarkResult]:
-    """
-    Run model parameter sweep for given backends and batch specs.
+    """Run model parameter sweep for given backends and batch specs.

    Args:
        backends: List of backend names
@@ -131,6 +130,7 @@ def run_model_parameter_sweep(

    Returns:
        List of BenchmarkResult objects
+
    """
    all_results = []

@@ -304,8 +304,7 @@ def run_parameter_sweep(
    sweep: ParameterSweep,
    console: Console,
 ) -> list[BenchmarkResult]:
-    """
-    Run parameter sweep for given backends and batch specs.
+    """Run parameter sweep for given backends and batch specs.

    Args:
        backends: List of backend names
@@ -316,6 +315,7 @@ def run_parameter_sweep(

    Returns:
        List of BenchmarkResult objects
+
    """
    all_results = []

@@ -395,8 +395,7 @@ def load_config_from_yaml(config_path: str) -> dict:


 def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
-    """
-    Generate batch specs from range specifications.
+    """Generate batch specs from range specifications.

    Args:
        ranges: List of range specifications, each containing:
@@ -420,6 +419,7 @@ def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
            }
        ]
        Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"]
+
    """
    all_specs = []

@@ -21,8 +21,7 @@ from vllm.triton_utils import triton


 def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
-    """
-    Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).
+    """Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).

    This ensures results are sorted by batch size first, then query length,
    then sequence length, rather than alphabetically.
@@ -354,13 +353,13 @@ class ResultsFormatter:
        backends: list[str],
        compare_to_fastest: bool = True,
    ):
-        """
-        Print results as a rich table.
+        """Print results as a rich table.

        Args:
            results: List of BenchmarkResult
            backends: List of backend names being compared
            compare_to_fastest: Show percentage comparison to fastest
+
        """
        # Group by batch spec, preserving first-occurrence order
        by_spec = {}
@@ -489,14 +488,14 @@ class ResultsFormatter:


 def setup_mla_dims(model_name: str = "deepseek-v3") -> dict:
-    """
-    Get MLA dimensions for known models.
+    """Get MLA dimensions for known models.

    Args:
        model_name: Model identifier

    Returns:
        Dict with MLA dimension configuration
+
    """
    configs = {
        "deepseek-v2": {
@@ -542,8 +541,7 @@ def get_attention_scale(head_dim: int) -> float:


 def is_mla_backend(backend: str) -> bool:
-    """
-    Check if backend is an MLA backend using the AttentionBackendEnum.
+    """Check if backend is an MLA backend using the AttentionBackendEnum.

    Args:
        backend: Backend name matching AttentionBackendEnum exactly
@@ -551,6 +549,7 @@ def is_mla_backend(backend: str) -> bool:

    Returns:
        True if the backend is an MLA backend, False otherwise
+
    """
    from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-MLA benchmark runner - shared utilities for MLA benchmarks.
+"""MLA benchmark runner - shared utilities for MLA benchmarks.

 This module provides helpers for running MLA backends without
 needing full VllmConfig integration.
@@ -93,8 +92,7 @@ def _ensure_single_rank_model_parallel() -> None:


 def _add_mock_methods_to_model_config(model_config: ModelConfig) -> None:
-    """
-    Add mock methods for layer-specific queries to ModelConfig.
+    """Add mock methods for layer-specific queries to ModelConfig.

    These methods are needed by metadata builders but aren't normally
    present on ModelConfig when used in benchmark contexts.
@@ -125,14 +123,15 @@ def create_minimal_vllm_config(
    kv_cache_dtype: str = "auto",
    sparse_mla_force_mqa: bool = False,
 ) -> VllmConfig:
-    """
-    Create minimal VllmConfig for MLA benchmarks.
+    """Create minimal VllmConfig for MLA benchmarks.

    Args:
        model_name: Model name (deepseek-v2, deepseek-v3, etc.) - used if mla_dims not
                    provided
        block_size: KV cache block size
        max_num_seqs: Maximum number of sequences
+        max_num_batched_tokens: Maximum number of batched tokens
+        max_model_len: Maximum model context length
        mla_dims: Optional custom MLA dimensions dict. If not provided, uses
                  setup_mla_dims(model_name)
        index_topk: Optional topk value for sparse MLA backends. If provided,
@@ -140,11 +139,13 @@ def create_minimal_vllm_config(
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
                        "trtllm"). Configures the attention config to force
                        the specified prefill backend.
+        kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8"
        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
                    forward_mqa (even prefill tokens).

    Returns:
        VllmConfig for benchmarking
+
    """
    # Get MLA dimensions - use provided or load from model name
    if mla_dims is None:
@@ -320,8 +321,7 @@ _BACKEND_PROPERTIES = {


 def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
+    """Get backend configuration from AttentionBackendEnum.

    Uses the registry to get the backend class and extract configuration
    from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.).
@@ -332,6 +332,7 @@ def _get_backend_config(backend: str) -> dict:

    Returns:
        Dict with backend configuration
+
    """
    from vllm.v1.attention.backend import MultipleOf
    from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -381,8 +382,7 @@ def _build_attention_metadata(
    device: torch.device,
    builder_instance,
 ) -> tuple:
-    """
-    Build attention metadata from batch requests.
+    """Build attention metadata from batch requests.

    Args:
        requests: List of BatchRequest objects
@@ -392,6 +392,7 @@ def _build_attention_metadata(

    Returns:
        Tuple of (metadata, kv_cache_num_blocks)
+
    """
    q_lens = [r.q_len for r in requests]
    kv_lens = [r.kv_len for r in requests]
@@ -472,8 +473,7 @@ def _create_input_tensors(
    device: torch.device,
    dtype: torch.dtype,
 ):
-    """
-    Create input tensors for both decode and prefill modes.
+    """Create input tensors for both decode and prefill modes.

    MLA requires different tensor formats for decode vs prefill:
    - Decode: Uses kv_lora_rank (512) dimension
@@ -490,6 +490,7 @@ def _create_input_tensors(
        Tuple of (decode_inputs, prefill_inputs)
        - decode_inputs: Query tensor(s) for decode mode
        - prefill_inputs: Dict with 'q', 'k_c_normed', 'k_pe', 'k_scale' for prefill
+
    """
    if query_format == "tuple":
        # Decode mode format: (q_nope, q_pe) where q_nope has kv_lora_rank dim
@@ -583,8 +584,7 @@ def _create_backend_impl(
    index_topk: int | None = None,
    kv_cache_dtype: str = "auto",
 ):
-    """
-    Create backend implementation instance.
+    """Create backend implementation instance.

    Args:
        backend_cfg: Backend configuration dict from _get_backend_config()
@@ -593,9 +593,11 @@ def _create_backend_impl(
        device: Target device
        max_num_tokens: Maximum number of tokens for sparse indexer buffer
        index_topk: Topk value for sparse MLA backends
+        kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8"

    Returns:
        Tuple of (impl, layer, builder_instance, indexer)
+
    """
    # Get classes from backend config (already resolved by _get_backend_config)
    impl_class = backend_cfg["impl_class"]
@@ -701,14 +703,14 @@ def _create_backend_impl(


 def _extract_mla_dims_from_config(config) -> dict | None:
-    """
-    Extract MLA dimensions from BenchmarkConfig if all required fields are present.
+    """Extract MLA dimensions from BenchmarkConfig if all required fields are present.

    Args:
        config: BenchmarkConfig instance

    Returns:
        Dict with MLA dimensions if all fields are provided, None otherwise
+
    """
    # Check if all MLA-specific fields are provided
    if all(
@@ -762,8 +764,7 @@ def _run_single_benchmark(
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> BenchmarkResult:
-    """
-    Run a single benchmark iteration.
+    """Run a single benchmark iteration.

    Args:
        config: BenchmarkConfig instance
@@ -774,6 +775,7 @@ def _run_single_benchmark(
        mla_dims: MLA dimension configuration
        device: Target device
        indexer: Optional MockIndexer for sparse backends
+        kv_cache_dtype: KV cache dtype. None keeps the backend default.
        output_scale: Static per-tensor FP8 scale for prefill output. None
            keeps the plain bf16 output (no quantization).
        fuse_quant_op: With output_scale set, True lets the prefill kernel write
@@ -782,6 +784,7 @@ def _run_single_benchmark(

    Returns:
        BenchmarkResult with timing statistics
+
    """
    # Parse batch spec
    requests = parse_batch_spec(config.batch_spec)
@@ -1044,8 +1047,7 @@ def _run_mla_benchmark_batched(
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
-    """
-    Unified batched MLA benchmark runner for all backends.
+    """Unified batched MLA benchmark runner for all backends.

    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
               flashinfer_mla_sparse, flashmla_sparse
@@ -1061,11 +1063,16 @@ def _run_mla_benchmark_batched(
        index_topk: Topk value for sparse MLA backends (default 2048)
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
            When set, forces the specified FlashAttention version for prefill.
+        output_scale: Static per-tensor FP8 scale for prefill output. None
+            keeps the plain bf16 output (no quantization).
+        fuse_quant_op: With output_scale set, True lets the prefill kernel write
+            FP8 directly; False runs bf16 attention then a standalone quant.
        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
            forward_mqa (even prefill tokens).

    Returns:
        List of BenchmarkResult objects
+
    """
    if not configs_with_params:
        return []
@@ -1277,8 +1284,7 @@ def run_mla_benchmark(
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
-    """
-    Unified MLA benchmark runner for all backends.
+    """Unified MLA benchmark runner for all backends.

    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
               flashinfer_mla_sparse, flashmla_sparse
@@ -1303,6 +1309,7 @@ def run_mla_benchmark(

    Returns:
        BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
+
    """
    # Normalize to batched mode: (config, threshold, num_splits)
    if isinstance(config, list):
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Standard attention benchmark runner - shared utilities for non-MLA benchmarks.
+"""Standard attention benchmark runner - shared utilities for non-MLA benchmarks.

 This module provides helpers for running standard attention backends
 (FlashAttention, Triton, FlashInfer) with real vLLM integration.
@@ -57,8 +56,7 @@ from vllm.v1.kv_cache_interface import (


 def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
+    """Get backend configuration from AttentionBackendEnum.

    Args:
        backend: Backend name matching AttentionBackendEnum exactly
@@ -66,6 +64,7 @@ def _get_backend_config(backend: str) -> dict:

    Returns:
        Dict with backend_class
+
    """
    from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -409,6 +408,7 @@ def _run_single_benchmark(
    Returns:
        (timing_stats, mem_stats) where timing_stats is a dict with
        mean/std/min/max in seconds per layer.
+
    """
    total_q = q_list[0].shape[0]
    v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
@@ -458,8 +458,7 @@ def _run_single_benchmark(


 def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
-    """
-    Run standard attention benchmark with real kernels.
+    """Run standard attention benchmark with real kernels.

    Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER

@@ -468,6 +467,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:

    Returns:
        BenchmarkResult with timing and memory statistics
+
    """
    device = torch.device(config.device)
    torch.accelerator.set_device_index(device)
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.
+"""Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.

 This benchmark runs the same workload twice:
 1. With VLLM_BATCH_INVARIANT=0 (baseline)
@@ -96,8 +95,7 @@ def run_benchmark_with_batch_invariant(
    batch_invariant: bool,
    seed: int = 12345,
 ) -> dict:
-    """
-    Run the benchmark with the specified configuration.
+    """Run the benchmark with the specified configuration.

    Returns a dict with timing and throughput metrics.
    """
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Micro benchmark comparing built-in hash(), SHA-256, and xxHash.
+"""Micro benchmark comparing built-in hash(), SHA-256, and xxHash.

 This focuses on a single test payload shaped like the prefix-cache hash input:
    (32-byte bytes object, 32-int tuple)
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark hidden state extraction throughput.
+"""Benchmark hidden state extraction throughput.

 Measures two modes:
  1. Baseline: bulk inference with max_tokens=1, no extraction.
@@ -137,7 +136,8 @@ async def _client_loop(
    client_id: int,
 ):
    """A single async client: pulls prompts, submits to engine, consumes
-    hidden states as soon as each request finishes."""
+    hidden states as soon as each request finishes.
+    """
    loop = asyncio.get_event_loop()
    while True:
        item = await prompt_queue.get()
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Offline benchmark to test the long document QA throughput.
+"""Offline benchmark to test the long document QA throughput.

 Example usage:
    # This workload samples 8 different prompts with a default input
@@ -49,14 +48,14 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser


 def test_long_document_qa(llm=None, sampling_params=None, prompts=None):
-    """
-    Test long document QA with the given prompts and sampling parameters.
+    """Test long document QA with the given prompts and sampling parameters.
    Print the time spent in processing all the prompts.

    Args:
        llm: The language model used for generating responses.
        sampling_params: Sampling parameter used to generate the response.
        prompts: A list of prompt strings to be processed by the LLM.
+
    """
    start_time = time.time()
    llm.generate(prompts, sampling_params=sampling_params)
@@ -65,8 +64,7 @@ def test_long_document_qa(llm=None, sampling_params=None, prompts=None):


 def repeat_prompts(prompts, repeat_count, mode: str):
-    """
-    Repeat each prompt in the list for a specified number of times.
+    """Repeat each prompt in the list for a specified number of times.
    The order of prompts in the output list depends on the mode.

    Args:
@@ -84,6 +82,7 @@ def repeat_prompts(prompts, repeat_count, mode: str):

    Raises:
        ValueError: If an invalid mode is provided.
+
    """
    print("Repeat mode: ", mode)
    if mode == "random":
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Simple benchmark to compare prefix-cache block hashing algorithms.
+"""Simple benchmark to compare prefix-cache block hashing algorithms.

 Example:
    python benchmark_prefix_block_hash.py --num-blocks 20000 --block-size 32
+
 """

 from __future__ import annotations
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the efficiency of prefix caching.
+"""Benchmark the efficiency of prefix caching.

 This script allows you to benchmark the performance of
 a model with and without prefix caching using either fixed prompts
@@ -102,6 +102,7 @@ class SampleRequest:
            images).
        prompt_len: The length of the prompt in tokens.
        expected_output_len: The expected length of the output in tokens.
+
    """

    prompt: str
@@ -305,8 +306,7 @@ async def get_request(
    request_rate: float,
    burstiness: float = 1.0,
 ) -> AsyncGenerator[tuple[int, SampleRequest], None]:
-    """
-    Asynchronously generates requests at a specified rate
+    """Asynchronously generates requests at a specified rate
    with OPTIONAL burstiness.

    Args:
@@ -322,6 +322,7 @@ async def get_request(
            A lower burstiness value (0 < burstiness < 1) results
            in more bursty requests, while a higher burstiness value
            (burstiness > 1) results in a more uniform arrival of requests.
+
    """
    input_requests = iter(input_requests)

@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
+"""Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.

 Compares:
 - apply_top_k_top_p_triton (Triton binary search)
@@ -47,8 +46,7 @@ def calculate_ops_pct(
    vocab_size: int,
    batch_size: int,
 ) -> float:
-    """
-    Calculate the percentage of active top-k and top-p operations.
+    """Calculate the percentage of active top-k and top-p operations.

    Returns percentage where 100% = batch_size ops.
    E.g., if all rows have both top-k and top-p active, returns 200%.
@@ -117,8 +115,7 @@ def benchmark_function(
    warmup_iters: int = 5,
    benchmark_iters: int = 20,
 ) -> tuple[float, int]:
-    """
-    Benchmark a function and return (avg_time_ms, peak_memory_bytes).
+    """Benchmark a function and return (avg_time_ms, peak_memory_bytes).

    Returns average time in milliseconds and peak memory usage.
    """
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark: Fused FP8 output quantization in merge_attn_states
+"""Benchmark: Fused FP8 output quantization in merge_attn_states.

 Compares fused vs unfused approaches for producing FP8-quantized merged
 attention output:
@@ -85,7 +84,8 @@ def make_inputs(

 def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes):
    """Build (num_tokens, num_heads, head_size, dtype_str) config tuples,
-    applying TP division to num_heads and skipping invalid combos."""
+    applying TP division to num_heads and skipping invalid combos.
+    """
    configs = []
    for (_, nh, hs), nt, dtype, tp in itertools.product(
        head_configs, num_tokens_list, input_dtypes, tp_sizes
@@ -24,9 +24,9 @@ def make_inputs(num_tokens, dtype):
    """Create inputs matching the real code path.

    Args:
-        contiguous_nope: If False, simulate the transposed BMM output
-                         (non-contiguous nope with stride pattern from
-                         [N,B,L].transpose(0,1)).
+        num_tokens: Number of tokens to generate inputs for.
+        dtype: dtype of the generated tensors.
+
    """
    # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L]
    raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda")
@@ -106,8 +106,7 @@ class BenchmarkTensors:


 def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool):
-    """
-    Reference triton quant kernel from,
+    """Reference triton quant kernel from,
    vllm.model_executor.layers.quantization.utils.fp8_utils
    """
    assert quant_out.size() == x.size()
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
+"""Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
 kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
 but use different quantization strategies and backends.
 """
@@ -190,7 +189,7 @@ def bench_run(
    torch.accelerator.synchronize()

    def bench_cuda_graph(graph, num_warmup=5, num_iters=100):
-        """Benchmark CUDA graph using events like benchmark_moe.py"""
+        """Benchmark CUDA graph using events like benchmark_moe.py."""
        # Warmup
        for _ in range(num_warmup):
            graph.replay()
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe
+"""Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe
 kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit
 activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8)
 and 16-bit activations.
@@ -2,8 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Benchmark script for device communicators:
+"""Benchmark script for device communicators:
 CustomAllreduce (oneshot, twoshot), PyNcclCommunicator,
 and SymmMemCommunicator (multimem, two-shot).

@@ -17,6 +16,7 @@ Usage:
 Example:
    torchrun --nproc_per_node=2 benchmark_device_communicators.py
    --sequence-lengths 512 1024 2048 --num-warmup 10 --num-trials 100
+
 """

 import json
@@ -241,7 +241,6 @@ class CommunicatorBenchmark:
        self, sequence_length: int, num_warmup: int, num_trials: int
    ) -> dict[str, float]:
        """Benchmark allreduce operations for all available communicators."""
-
        results = {}

        # Define communicators with their benchmark functions
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Benchmark for FlashInfer fused collective operations vs standard operations.
+"""Benchmark for FlashInfer fused collective operations vs standard operations.

 This benchmark compares:
 1. FlashInfer's allreduce_fusion with trtllm backend
@@ -446,9 +445,15 @@ def run_benchmarks(
    """Run all benchmarks for given configuration.

    Args:
+        num_tokens: Number of tokens in the benchmark input.
+        hidden_dim: Hidden dimension of the benchmark input.
+        dtype: dtype of the benchmark input.
+        use_residual: Whether to include a residual add in the fused op.
+        no_oneshot: Skip the one-shot allreduce variants.
        allreduce_params: Shared parameters for FlashInfer fused allreduce.
        workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
        quant_modes: Set of quantization modes: "none", "fp8", "fp4".
+
    """
    (
        input_tensor,
@@ -10,6 +10,7 @@ Example:

    .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \
      --output /tmp/attn_res.json
+
 """

 from __future__ import annotations
@@ -16,6 +16,7 @@ state traffic itself.
 Example:
    python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \
        --tokens 1 8 32 64 128 --heads 12
+
 """

 import argparse
@@ -27,6 +27,7 @@ Examples:

 For multi-node runs, launch one ``torchrun`` agent per node and use a shared
 rendezvous endpoint.
+
 """

 from __future__ import annotations
@@ -102,9 +102,7 @@ def make_rand_tensors(
    num_slices: int,
    device: str = "cuda",
 ) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]:
-    """
-    Make LoRA input/output matrices.
-    """
+    """Make LoRA input/output matrices."""
    A = torch.rand(a_shape, dtype=a_dtype).to(device)

    # LoRA weights column major
@@ -117,8 +115,7 @@ def make_rand_tensors(
 def make_prompt_lora_mapping(
    num_prompts: int, num_active_loras: int, sort_by_lora_id: bool, device: str
 ) -> torch.Tensor:
-    """
-    All prompts are mapped to a LoRA ID in range [0, num_active_loras).
+    """All prompts are mapped to a LoRA ID in range [0, num_active_loras).
    where 0 refers to first lora, 1 refers to second lora and so on.
    """
    assert num_active_loras > 0
@@ -147,9 +144,7 @@ def make_token_lora_mapping(
    seq_len_tensor: torch.Tensor,
    device: str,
 ):
-    """
-    Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor
-    """
+    """Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor."""
    assert prompt_lora_mapping.shape[0] == num_prompts

    # token to lora index mapping
@@ -174,8 +169,7 @@ def ref_group_gemm(
    scaling: float,
    add_inputs: bool | None,
 ):
-    """
-    Torch group gemm reference implementation to test correctness of
+    """Torch group gemm reference implementation to test correctness of
    benchmarking operations.
    """
    batches = seq_lens_cpu.size(0)
@@ -198,9 +192,7 @@ def ref_group_gemm(


 class OpType(Enum):
-    """
-    LoRA Ops to benchmark and its properties.
-    """
+    """LoRA Ops to benchmark and its properties."""

    LORA_SHRINK = auto()
    LORA_EXPAND = auto()
@@ -291,9 +283,7 @@ class OpType(Enum):
    def matmul_dtypes(
        self, op_dtype: torch.dtype
    ) -> tuple[torch.dtype, torch.dtype, torch.dtype]:
-        """
-        return a type, b type and c type for A x B = C
-        """
+        """Return a type, b type and c type for A x B = C."""
        if self.is_shrink_fn():
            return op_dtype, op_dtype, torch.float32
        elif self.is_expand_fn():
@@ -338,8 +328,7 @@ class OpType(Enum):
        top_k_num: int | None = None,
        num_experts: int | None = None,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        """
-        Given num_slices, return the shapes of the A, B, and C matrices
+        """Given num_slices, return the shapes of the A, B, and C matrices
        in A x B = C, for the op_type
        """
        m, k, n = self.mkn(batch_size, seq_length, hidden_size, lora_rank)
@@ -419,9 +408,7 @@ class OpType(Enum):

@dataclass
 class BenchmarkContext:
-    """
-    LoRA benchmark context
-    """
+    """LoRA benchmark context."""

    batch_size: int
    hidden_size: int
@@ -467,9 +454,7 @@ class BenchmarkContext:

@dataclass
 class BenchmarkTensors:
-    """
-    Input/Output tensors used for benchmarks
-    """
+    """Input/Output tensors used for benchmarks."""

    # matmul tensors
    input: torch.Tensor
@@ -555,9 +540,7 @@ class BenchmarkTensors:
        )

    def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None:
-        """
-        Fails asserts when non-conformality is detected.
-        """
+        """Fails asserts when non-conformality is detected."""
        num_tokens = (
            self.input.shape[1]
            if op_type.is_fused_moe_lora_expand_fn()
@@ -577,9 +560,7 @@ class BenchmarkTensors:
        )

    def to_device(self, device: str):
-        """
-        Transfer tensors to device if the tensors aren't already on the device
-        """
+        """Transfer tensors to device if the tensors aren't already on the device."""

        def to_device(tensor: torch.Tensor):
            if tensor.device != device:
@@ -604,9 +585,7 @@ class BenchmarkTensors:
            )

    def metadata(self, ctx: BenchmarkContext, op_type: OpType) -> tuple[int, int, int]:
-        """
-        Return num_seqs, num_tokens and max_seq_len
-        """
+        """Return num_seqs, num_tokens and max_seq_len."""
        num_seqs = self.seq_lens.shape[0]
        num_tokens = self.get_num_tokens(
            self.lora_kernel_meta.token_lora_mapping.shape[0], ctx.top_k_num, op_type
@@ -630,8 +609,7 @@ class BenchmarkTensors:
            expert_map: torch.Tensor | None = None,
            pad_sorted_ids: bool = False,
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            """
-            Aligns tokens and experts into block-sized chunks for LoRA-based
+            """Aligns tokens and experts into block-sized chunks for LoRA-based
            mixture-of-experts (MoE) execution.
            """
            max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
@@ -959,8 +937,7 @@ class BenchmarkTensors:
    def test_correctness(
        self, op_type: OpType, expand_fn_add_inputs: bool | None
    ) -> bool:
-        """
-        Test correctness of op_type implementation against a grouped gemm
+        """Test correctness of op_type implementation against a grouped gemm
        reference implementation.
        """
        seq_lens_cpu = self.seq_lens.to(device="cpu")
@@ -1070,15 +1047,13 @@ def bench_torch_mm(
    op_type: OpType,
    cuda_graph_nops: int | None = None,
 ) -> TMeasurement:
-    """
-    Benchmark basic torch.mm as a roofline.
+    """Benchmark basic torch.mm as a roofline.

    When all the input tokens have the same LoRA ID, the LoRA kernels are just
    a matmul. This torch.mm benchmark serves as a roofline for that case.

    input op_type is used in determining the m, k, n dimensions for the matmul.
    """
-
    batch_size, hidden_size, lora_rank, seq_length, dtype = (
        ctx.batch_size,
        ctx.hidden_size,
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
+"""Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
 in MLA (Multi-head Latent Attention) prefill.

 This validates that the optimization from commit 8d4142bd is beneficial across
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing old vs new default fused MoE configs.
+"""Benchmark comparing old vs new default fused MoE configs.

 Runs the triton fused_moe kernel with three configurations for each scenario:
  1. Tuned config (from JSON file, if available) — the target to match
@@ -33,7 +32,8 @@ FP8_DTYPE = current_platform.fp8_dtype()

 def old_default_config(M, E, N, K, topk, dtype=None, block_shape=None):
    """The original defaults before https://github.com/vllm-project/vllm/pull/34846,
-    for comparison."""
+    for comparison.
+    """
    if dtype == "fp8_w8a8" and block_shape is not None:
        return {
            "BLOCK_SIZE_M": 64,
@@ -15,7 +15,7 @@ from vllm.platforms import current_platform

@contextmanager
 def _triton_mode():
-    """Temporarily force the Triton fallback path"""
+    """Temporarily force the Triton fallback path."""
    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
        yield

@@ -16,7 +16,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE


 def with_triton_mode(fn):
-    """Temporarily force the Triton fallback path"""
+    """Temporarily force the Triton fallback path."""

    def wrapped(*args, **kwargs):
        with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
@@ -124,8 +124,7 @@ def compute_geomean_speedups(
    speedup_cols: list[str],
    groupby_cols: list[str] | None = None,
 ) -> pd.DataFrame:
-    """
-    Compute geometric mean speedups over a baseline column.
+    """Compute geometric mean speedups over a baseline column.

    Args:
        df: Input dataframe
@@ -135,6 +134,7 @@ def compute_geomean_speedups(

    Returns:
        pd.DataFrame with geometric mean speedups
+
    """
    from scipy.stats import gmean

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
+"""Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.

 Usage:
    python benchmark_int4_gemm.py
@@ -32,7 +32,6 @@ def run_benchmark(
    device: str = "cuda",
 ) -> float:
    """Return latency (seconds) for given num_tokens."""
-
    if kv_cache_dtype == "fp8" and head_size % 16:
        raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")

@@ -37,7 +37,6 @@ def run_benchmark(
    device: str = "cuda",
 ) -> float:
    """Return latency (seconds) for given num_tokens."""
-
    if kv_cache_dtype == "fp8" and head_size % 16:
        raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")

@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark and tuning script for the Mamba selective_state_update kernel.
+"""Benchmark and tuning script for the Mamba selective_state_update kernel.

 Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
 an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
@@ -124,8 +123,7 @@ def benchmark_config(
    num_warmup: int = 20,
    graph_batch_size: int = 10,
 ) -> float | None:
-    """
-    Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
+    """Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
    Returns elapsed time in microseconds, or None on error.

    Uses accelerator graph capture-and-replay to isolate kernel time from
@@ -356,8 +354,7 @@ def validate_configs(
    rtol: float = 1e-2,
    state_dtype: torch.dtype | None = None,
 ) -> dict[int, bool]:
-    """
-    For every (effective_batch, batch, nheads) in *active* that has a tuned
+    """For every (effective_batch, batch, nheads) in *active* that has a tuned
    config, run the kernel with that config and compare against the reference.
    Returns {effective_batch: passed}.
    """
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""
-Comprehensive 3-way SiLU Benchmark Suite
+"""Comprehensive 3-way SiLU Benchmark Suite.

 This benchmark compares three SiLU implementations:
 1. SiLU V2 (CUDA) - Optimized CUDA kernel implementation
@@ -110,7 +109,7 @@ def silu_mul_fp8_quant_deep_gemm_triton(
    eps: float = 1e-10,
    expert_offsets: torch.Tensor = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
+    """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales.

    y has shape (E, T, 2*H). The first half of the last dimension is
    silu-activated, multiplied by the second half, then quantized into FP8.
@@ -206,7 +205,7 @@ def benchmark(
    iterations_per_run: int = 20,
 ):
    def generate_data(seed_offset=0):
-        """Generate input data with given seed offset"""
+        """Generate input data with given seed offset."""
        set_random_seed(42 + seed_offset)
        y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()

@@ -10,6 +10,7 @@ with a static scale file.
 Example:
    python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \
        --seq-lens 2304 4096 8192 16384 --head-dim 72
+
 """

 from types import SimpleNamespace
@@ -97,7 +97,8 @@ def graph_replay(fn):

 def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens):
    """The fused kernel must be bit-exact with the path it replaces, on every
-    rank of the shard."""
+    rank of the shard.
+    """
    dtype = torch.bfloat16
    for tp_rank in range(tp_size):
        input_ids, weight, shard_indices = make_inputs(
@@ -57,10 +57,12 @@ def w8a8_block_matmul(
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization.
                    It should be 2-dim, e.g., [128, 128].
+        config: Triton kernel launch config, e.g. tile sizes and num_warps.
        output_dtype: The dtype of the returned tensor.

    Returns:
        torch.Tensor: The result of matmul.
+
    """
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generic benchmark harness for vLLM IR ops.
+"""Generic benchmark harness for vLLM IR ops.

 Usage:
    python benchmarks/kernels/ir/bench_ir_ops.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shape configurations for IR op benchmarks.
-"""
+"""Shape configurations for IR op benchmarks."""

 import torch

@@ -17,8 +17,7 @@ class CudaGraphBenchParams:

@dataclasses.dataclass
 class ArgPool:
-    """
-    When some argument of the benchmarking function is annotated with this type,
+    """When some argument of the benchmarking function is annotated with this type,
    the benchmarking class (BenchMM) will collapse the argument to a pick a
    single value from the given list of values, during function invocation.
    For every invocation during a benchmarking run, it will choose a
@@ -148,8 +148,7 @@ class LognormalDistribution(Distribution):
    def _generate_lognormal_by_median(
        target_average: int, median_ratio: float
    ) -> tuple[float, float]:
-        """
-        Compute (mu, sigma) for a lognormal distribution given:
+        """Compute (mu, sigma) for a lognormal distribution given:
        - a target average (mean of the distribution)
        - a ratio of median / mean (controls skewness), assume mean > median

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Download dataset from:
+"""Download dataset from:
 https://huggingface.co/datasets/philschmid/sharegpt-raw/blob/main/sharegpt_20230401_clean_lang_split.json

 Convert to OpenAI API:
@@ -15,6 +15,7 @@ Examples:
    python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
    python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
        --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
+
 """

 import argparse