iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add pydocstyle to the ruff rules (#52136)

Browse Source 
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>
This commit is contained in:
Harry Mellor
2026-09-16 16:44:29 +00:00
committed by GitHub
co-authored by Claude Opus 5
parent 804e5377aa
commit 8c1557a79c
2002 changed files with 15024 additions and 14269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
LM eval harness on model to compare vs HF baseline computed offline.
"""LM eval harness on model to compare vs HF baseline computed offline.
Configs are found in configs/$MODEL.yaml
pytest -s -v test_lm_eval_correctness.py \
.buildkite/performance-benchmarks/scripts/compare-json-results.py
+2 -4
View File
@@ -83,8 +83,7 @@ def compare_data_columns(
drop_column: str,
debug: bool = False,
):
"""
Align concatenation by keys derived from info_cols instead of row order.
"""Align concatenation by keys derived from info_cols instead of row order.
- Pick one canonical key list: subset of info_cols present in ALL files.
- For each file: set index to those keys, aggregate duplicates
(mean for metric, first for names).
@@ -363,8 +362,7 @@ def _apply_two_decimals(
# Export helpers (Excel + CSV)
# -----------------------------
def _sanitize_sheet_name(name: str) -> str:
"""
Excel sheet constraints:
"""Excel sheet constraints:
- max 31 chars
- cannot contain: : \ / ? * [ ]
- cannot be empty
.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py
+1 -2
View File
@@ -93,8 +93,7 @@ def results_to_json(latency, throughput, serving):
def get_size_with_unit(bytes, suffix="B"):
"""
Scale bytes to its proper format
"""Scale bytes to its proper format
e.g:
1253656 => '1.20MB'
1253656678 => '1.17GB'
.buildkite/scripts/generate-nightly-index.py
+8 -13
View File
@@ -18,8 +18,7 @@ import regex as re
def normalize_package_name(name: str) -> str:
"""
Normalize package name according to PEP 503.
"""Normalize package name according to PEP 503.
https://peps.python.org/pep-0503/#normalized-names
Replace runs of underscores, hyphens, and periods with a single hyphen,
@@ -55,17 +54,18 @@ class WheelFileInfo:
def parse_from_filename(file: str) -> WheelFileInfo:
"""
Parse wheel file name to extract metadata.
"""Parse wheel file name to extract metadata.
The format of wheel names:
{package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
All versions could contain a variant like '+cu129' or '.cpu' or `.rocm` (or not).
Example:
vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl
vllm-0.11.1rc8.dev14+gaa384b3c0-cp38-abi3-manylinux2014_aarch64.whl
vllm-0.11.1rc8.dev14+gaa384b3c0.cu130-cp38-abi3-manylinux1_x86_64.whl
"""
wheel_file_re = re.compile(
r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
@@ -111,9 +111,7 @@ def parse_from_filename(file: str) -> WheelFileInfo:
def generate_project_list(subdir_names: list[str], comment: str = "") -> str:
"""
Generate project list HTML content linking to each project & variant subdirectory.
"""
"""Generate project list HTML content linking to each project & variant subdirectory."""
href_tags = []
for name in sorted(subdir_names):
name = name.strip("/").strip(".")
@@ -127,9 +125,7 @@ def generate_package_index_and_metadata(
index_base_dir: Path,
comment: str = "",
) -> tuple[str, str]:
"""
Generate package index HTML content for a specific package, linking to actual wheel files.
"""
"""Generate package index HTML content for a specific package, linking to actual wheel files."""
href_tags = []
metadata = []
for file in sorted(wheel_files, key=lambda x: x.filename):
@@ -156,8 +152,7 @@ def generate_index_and_metadata(
alias_to_default: str | None = None,
comment: str = "",
):
"""
Generate index for all wheel files.
"""Generate index for all wheel files.
Args:
whl_files (list[str]): List of wheel files (must be directly under `wheel_base_dir`).
@@ -216,8 +211,8 @@ def generate_index_and_metadata(
},
...
]
"""
"""
parsed_files = [parse_from_filename(f) for f in whl_files]
if not parsed_files:
benchmarks/attention_benchmarks/batch_spec.py
+14 -14
View File
@@ -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)
benchmarks/attention_benchmarks/benchmark.py
+10 -10
View File
@@ -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 = []
benchmarks/attention_benchmarks/common.py
+7 -8
View File
@@ -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
benchmarks/attention_benchmarks/mla_runner.py
+29 -22
View File
@@ -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):
benchmarks/attention_benchmarks/runner.py
+6 -6
View File
@@ -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)
benchmarks/benchmark_batch_invariance.py
+2 -4
View File
@@ -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.
"""
benchmarks/benchmark_hash.py
+1 -2
View File
@@ -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)
benchmarks/benchmark_hidden_state_extraction.py
+3 -3
View File
@@ -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()
benchmarks/benchmark_long_document_qa_throughput.py
+5 -6
View File
@@ -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":
benchmarks/benchmark_prefix_block_hash.py
+2 -2
View File
@@ -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
benchmarks/benchmark_prefix_caching.py
+1 -2
View File
@@ -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
benchmarks/benchmark_serving_structured_output.py
+3 -2
View File
@@ -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)
benchmarks/benchmark_topk_topp.py
+3 -6
View File
@@ -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.
"""
benchmarks/fused_kernels/merge_attn_states_benchmarks.py
+3 -3
View File
@@ -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
benchmarks/kernels/bench_concat_mla_q.py
+3 -3
View File
@@ -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")
benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py
+1 -2
View File
@@ -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()
benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+2 -3
View File
@@ -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()
benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py
+1 -2
View File
@@ -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.
benchmarks/kernels/benchmark_device_communicators.py
+2 -3
View File
@@ -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
benchmarks/kernels/benchmark_fused_collective.py
+7 -2
View File
@@ -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,
benchmarks/kernels/benchmark_kimi_k3_attn_res.py
+1
View File
@@ -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
benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
+1
View File
@@ -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
benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
+1
View File
@@ -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
benchmarks/kernels/benchmark_lora.py
+15 -40
View File
@@ -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,
benchmarks/kernels/benchmark_mla_k_concat.py
+1 -2
View File
@@ -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
benchmarks/kernels/benchmark_moe_defaults.py
+3 -3
View File
@@ -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,
benchmarks/kernels/benchmark_per_token_group_quant.py
+1 -1
View File
@@ -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
benchmarks/kernels/benchmark_per_token_quant_fp8.py
+3 -3
View File
@@ -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
benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
+1 -2
View File
@@ -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
benchmarks/kernels/benchmark_reshape_and_cache.py
-1
View File
@@ -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.")
benchmarks/kernels/benchmark_reshape_and_cache_flash.py
-1
View File
@@ -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.")
benchmarks/kernels/benchmark_selective_state_update.py
+3 -6
View File
@@ -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}.
"""
benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
+3 -4
View File
@@ -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()
benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py
+1
View File
@@ -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
benchmarks/kernels/benchmark_vocab_parallel_embedding.py
+2 -1
View File
@@ -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(
benchmarks/kernels/benchmark_w8a8_block_fp8.py
+2
View File
@@ -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]
benchmarks/kernels/ir/bench_ir_ops.py
+1 -2
View File
@@ -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
benchmarks/kernels/ir/shapes.py
+1 -3
View File
@@ -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
benchmarks/kernels/utils.py
+1 -2
View File
@@ -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
benchmarks/multi_turn/bench_dataset.py
+1 -2
View File
@@ -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
benchmarks/multi_turn/convert_sharegpt_to_openai.py
+1 -2
View File
@@ -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:
benchmarks/replayssm/e2e_decode_speedup.py
+1
View File
@@ -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
csrc/cpu/generate_cpu_attn_dispatch.py
+1 -3
View File
@@ -1,9 +1,7 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Generate CPU attention dispatch switch cases and kernel instantiations.
"""
"""Generate CPU attention dispatch switch cases and kernel instantiations."""
import os
docs/mkdocs/gen_files/generate_argparse.py
+3 -3
View File
@@ -98,8 +98,7 @@ mock_if_no_torch("torch", PydanticMagicMock(name="torch", nn=mock_nn))
# Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__
# We need to return the corresponding number of inputs, as IR infra will assert it
def get_outputs(native_fn: Callable) -> str:
"""
Extract output schema from function's return type annotation,
"""Extract output schema from function's return type annotation,
e.g. 'Tensor' or 'Tensor, Tensor'.
"""
import typing
@@ -249,11 +248,12 @@ 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
add_cli_args: The `add_cli_args` classmethod to build the parser from.
**kwargs: Additional keyword arguments to pass to `cls.add_cli_args`.
Returns:
FlexibleArgumentParser: A parser with markdown formatting for the class.
"""
try:
parser = FlexibleArgumentParser(add_json_tip=False)
docs/mkdocs/gen_files/generate_attention_backends.py
+4 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Generates documentation table for attention backends showing feature support.
"""Generates documentation table for attention backends showing feature support.
This script parses all registered attention backends using AST (no imports needed)
and generates a markdown table showing what features each backend supports,
@@ -304,6 +303,7 @@ def parse_mla_prefill_registry() -> dict[str, str]:
Returns:
A dict mapping backend names to their class paths.
"""
if not MLA_PREFILL_REGISTRY_FILE.exists():
return {}
@@ -325,6 +325,7 @@ def parse_mla_prefill_priorities() -> dict[str, list[str]]:
Returns:
A dict with keys like 'blackwell' and 'default' containing
lists of backend enum names in priority order.
"""
if not MLA_PREFILL_SELECTOR_FILE.exists():
return {}
@@ -549,6 +550,7 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
Returns:
A dict with backend properties, or None if parsing fails.
"""
file_path = get_file_from_class_path(class_path)
if file_path is None:
docs/mkdocs/gen_files/generate_examples.py
+8 -6
View File
@@ -49,8 +49,7 @@ def title(text: str) -> str:
@dataclass
class Example:
"""
Example class for generating documentation content from a given path.
"""Example class for generating documentation content from a given path.
Attributes:
path (Path): The path to the main directory or file.
@@ -64,6 +63,7 @@ class Example:
Methods:
generate() -> str: Generates the documentation content.
"""
path: Path
@@ -75,7 +75,8 @@ class Example:
If path is a file, it returns the path itself. If path is a directory, it
searches for Markdown files (*.md) in the directory and returns the first one
found. If no Markdown files are found, it returns None."""
found. If no Markdown files are found, it returns None.
"""
# Single file example
if self.path.is_file():
return self.path
@@ -90,7 +91,8 @@ class Example:
"""Determine other files in the directory excluding the main file.
If path is a file, it returns an empty list. Otherwise, it returns every file
in the directory except the main file in a list."""
in the directory except the main file in a list.
"""
# Single file example
if self.path.is_file():
return []
@@ -116,8 +118,7 @@ class Example:
raise ValueError(f"Title not found in {self.main_file}")
def fix_relative_links(self, content: str) -> str:
"""
Fix relative links in markdown content by converting them to gh-file
"""Fix relative links in markdown content by converting them to gh-file
format.
Args:
@@ -125,6 +126,7 @@ class Example:
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 #
docs/mkdocs/gen_files/generated_content.py
+1
View File
@@ -35,6 +35,7 @@ def fill_markers(doc_path: str, blocks: dict[str, str]) -> None:
FileNotFoundError: If the source page does not exist.
ValueError: If an expected marker is missing, or any `gen:` marker is
left unsubstituted after filling.
"""
source = DOCS_DIR / doc_path
if not source.exists():
docs/mkdocs/hooks/autoref_code.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
MkDocs hook to automatically convert inline code references to API doc links.
"""MkDocs hook to automatically convert inline code references to API doc links.
For example, `WeightTransferConfig` becomes
[`WeightTransferConfig`][vllm.config.WeightTransferConfig]
docs/mkdocs/hooks/url_schemes.py
+3 -8
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
MkDocs hook + markdown extension to enable the following links to render correctly,
"""MkDocs hook + markdown extension to enable the following links to render correctly,
including inside content included via pymdownx.snippets:
- Relative file links outside of the `docs/` directory, e.g.:
@@ -71,9 +70,7 @@ class UrlSchemesPreprocessor(Preprocessor):
return lines
def replace_relative_link(match: re.Match) -> str:
"""
Replace relative file links with URLs if they point outside the docs dir.
"""
"""Replace relative file links pointing outside the docs dir with URLs."""
title = match.group("title")
path = match.group("path")
path = ((DOC_DIR / page.file.src_uri).parent / path).resolve()
@@ -91,9 +88,7 @@ class UrlSchemesPreprocessor(Preprocessor):
return f"[{gh_icon} {title}]({url})"
def replace_github_link(match: re.Match) -> str:
"""
Replace GitHub issue, PR, and project links with enhanced Markdown links.
"""
"""Replace GitHub issue, PR, and project links with rich Markdown links."""
repo = match.group("repo")
type = match.group("type")
number = match.group("number")
examples/applications/api_server/server.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
NOTE: This API server is used only for demonstrating usage of AsyncEngine
"""NOTE: This API server is used only for demonstrating usage of AsyncEngine
and simple performance benchmarks. It is not intended for production use.
For production use, we recommend using our OpenAI compatible server.
We are also not going to accept PRs modifying this file.
examples/applications/chatbot/streamlit_openai_chatbot_webserver.py
+5 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
vLLM Chat Assistant - A Streamlit Web Interface
"""vLLM Chat Assistant - A Streamlit Web Interface.
A streamlined chat interface that quickly integrates
with vLLM API server.
@@ -78,6 +77,7 @@ def create_new_chat_session():
- current_session: Sets to new session ID
- active_session: Sets to new session ID
- messages: Resets to empty list
"""
session_id = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
st.session_state.sessions[session_id] = []
@@ -101,6 +101,7 @@ def switch_to_chat_session(session_id):
- current_session: Updated to specified session_id
- active_session: Updated to specified session_id
- messages: Loaded from sessions[session_id]
"""
st.session_state.current_session = session_id
st.session_state.active_session = session_id
@@ -134,6 +135,7 @@ def get_llm_response(messages, model, reason, content_ph=None, reasoning_ph=None
Note:
The function uses streamlit placeholders for live updates.
When reason=True, the reasoning process appears above the response.
"""
full_text = ""
think_text = ""
@@ -256,6 +258,7 @@ def server_supports_reasoning():
Returns:
bool: True if the model supports reasoning, False otherwise
"""
resp = client.chat.completions.create(
model=model,
examples/applications/rag/retrieval_augmented_generation_with_langchain.py
+10 -27
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Retrieval Augmented Generation (RAG) Implementation with Langchain
"""Retrieval Augmented Generation (RAG) Implementation with Langchain
==================================================================
This script demonstrates a RAG implementation using LangChain, Milvus
@@ -35,6 +34,7 @@ Notes:
- Ensure both vLLM services are running before executing
- Default ports: 8000 (embedding), 8001 (chat)
- First run may take time to download models
"""
import argparse
@@ -52,9 +52,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
def load_and_split_documents(config: dict[str, Any]):
"""
Load and split documents from web URL
"""
"""Load and split documents from web URL."""
try:
loader = WebBaseLoader(web_paths=(config["url"],))
docs = loader.load()
@@ -70,9 +68,7 @@ def load_and_split_documents(config: dict[str, Any]):
def init_vectorstore(config: dict[str, Any], documents: list[Document]):
"""
Initialize vector store with documents
"""
"""Initialize vector store with documents."""
return Milvus.from_documents(
documents=documents,
embedding=OpenAIEmbeddings(
@@ -86,9 +82,7 @@ def init_vectorstore(config: dict[str, Any], documents: list[Document]):
def init_llm(config: dict[str, Any]):
"""
Initialize llm
"""
"""Initialize llm."""
return ChatOpenAI(
model=config["chat_model"],
openai_api_key=config["vllm_api_key"],
@@ -97,9 +91,7 @@ def init_llm(config: dict[str, Any]):
def get_qa_prompt():
"""
Get question answering prompt template
"""
"""Get question answering prompt template."""
template = """You are an assistant for question-answering tasks.
Use the following pieces of retrieved context to answer the question.
If you don't know the answer, just say that you don't know.
@@ -112,16 +104,12 @@ Answer:
def format_docs(docs: list[Document]):
"""
Format documents for prompt
"""
"""Format documents for prompt."""
return "\n\n".join(doc.page_content for doc in docs)
def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
"""
Set up question answering chain
"""
"""Set up question answering chain."""
return (
{
"context": retriever | format_docs,
@@ -134,9 +122,7 @@ def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
def get_parser() -> argparse.ArgumentParser:
"""
Parse command line arguments
"""
"""Parse command line arguments."""
parser = argparse.ArgumentParser(description="RAG with vLLM and langchain")
# Add command line arguments
@@ -192,10 +178,7 @@ def get_parser() -> argparse.ArgumentParser:
def init_config(args: Namespace):
"""
Initialize configuration settings from command line arguments
"""
"""Initialize configuration settings from command line arguments."""
return {
"vllm_api_key": args.vllm_api_key,
"vllm_embedding_endpoint": args.vllm_embedding_endpoint,
examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py
+9 -9
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
"""RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
================================================================
This script demonstrates a RAG system using:
@@ -35,6 +34,7 @@ Notes:
- Ensure both vLLM services are running before executing
- Default ports: 8000 (embedding), 8001 (chat)
- First run may take time to download models
"""
import argparse
@@ -50,7 +50,7 @@ from llama_index.vector_stores.milvus import MilvusVectorStore
def init_config(args: Namespace):
"""Initialize configuration with command line arguments"""
"""Initialize configuration with command line arguments."""
return {
"url": args.url,
"embedding_model": args.embedding_model,
@@ -66,12 +66,12 @@ def init_config(args: Namespace):
def load_documents(url: str) -> list:
"""Load and process web documents"""
"""Load and process web documents."""
return SimpleWebPageReader(html_to_text=True).load_data([url])
def setup_models(config: dict[str, Any]):
"""Configure embedding and chat models"""
"""Configure embedding and chat models."""
Settings.embed_model = OpenAILikeEmbedding(
api_base=config["embedding_endpoint"],
api_key=config["vllm_api_key"],
@@ -96,14 +96,14 @@ def setup_models(config: dict[str, Any]):
def setup_vector_store(db_path: str) -> MilvusVectorStore:
"""Initialize vector store"""
"""Initialize vector store."""
sample_emb = Settings.embed_model.get_text_embedding("test")
print(f"Embedding dimension: {len(sample_emb)}")
return MilvusVectorStore(uri=db_path, dim=len(sample_emb), overwrite=True)
def create_index(documents: list, vector_store: MilvusVectorStore):
"""Create document index"""
"""Create document index."""
storage_context = StorageContext.from_defaults(vector_store=vector_store)
return VectorStoreIndex.from_documents(
documents,
@@ -112,13 +112,13 @@ def create_index(documents: list, vector_store: MilvusVectorStore):
def query_document(index: VectorStoreIndex, question: str, top_k: int):
"""Query document with given question"""
"""Query document with given question."""
query_engine = index.as_query_engine(similarity_top_k=top_k)
return query_engine.query(question)
def get_parser() -> argparse.ArgumentParser:
"""Parse command line arguments"""
"""Parse command line arguments."""
parser = argparse.ArgumentParser(description="RAG with vLLM and LlamaIndex")
# Add command line arguments
examples/deployment/async_llm_streaming.py
+2 -4
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine).
"""Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine).
This script shows the core functionality of vLLM's AsyncLLM engine for streaming
token-by-token output in offline inference scenarios. It demonstrates DELTA mode
@@ -20,8 +19,7 @@ from vllm.v1.engine.async_llm import AsyncLLM
async def stream_response(engine: AsyncLLM, prompt: str, request_id: str) -> None:
"""
Stream response from AsyncLLM and display tokens as they arrive.
"""Stream response from AsyncLLM and display tokens as they arrive.
This function demonstrates the core streaming pattern:
1. Create SamplingParams with DELTA output kind
examples/deployment/llm_engine_example.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates using the `LLMEngine`
"""This file demonstrates using the `LLMEngine`
for processing prompts with various sampling parameters.
"""
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+7 -13
View File
@@ -1,8 +1,7 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
disagg_encoder_proxy.py
"""disagg_encoder_proxy.py
Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
clusters:
@@ -200,8 +199,7 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
def extract_mm_items(request_data: dict) -> list[dict]:
"""
Return *all* image/audio/video items that appear anywhere in `messages`.
"""Return *all* image/audio/video items that appear anywhere in `messages`.
Each returned dict looks like:
{ "type": "image_url", "image_url": {...} }
@@ -224,8 +222,7 @@ async def fanout_encoder_primer(
req_id: str,
consumer_zmq: str | None = None,
) -> tuple[dict[int, dict], dict[str, Any]]:
"""
1. Build one request *per MM item* with all text removed.
"""1. Build one request *per MM item* with all text removed.
2. Send them concurrently to the encode cluster.
3. Raise if any of them fails.
@@ -402,8 +399,7 @@ async def maybe_prefill(
p_url: str,
req_id: str,
) -> dict:
"""
- Do prefill-only task if p_url exist;
"""- Do prefill-only task if p_url exist;
- Return a new body carrying kv transfer params (for nixl connector)
- Else, skip and return the original request data for decode
@@ -849,14 +845,14 @@ async def _post_if_available(
payload: dict,
headers: dict,
) -> dict | None:
"""
POST `payload` to `url`.
"""POST `payload` to `url`.
Returns
-------
• The decoded JSON body on success (2xx)
• None if the endpoint does not exist (404)
• Raises for anything else.
"""
try:
resp = await session.post(url, json=payload, headers=headers)
@@ -877,9 +873,7 @@ async def _post_if_available(
async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str):
"""
Fire & forget to both clusters, tolerate 404.
"""
"""Fire & forget to both clusters, tolerate 404."""
headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
encode_task = _post_if_available(
examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file provides a disaggregated prefilling proxy demo to demonstrate an
"""This file provides a disaggregated prefilling proxy demo to demonstrate an
example usage of XpYd disaggregated prefilling.
We can launch multiple vllm instances (2 for prefill and 2 for decode), and
launch this proxy demo through:
examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py
+3 -3
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer
"""Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer.
This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment,
routing multi-turn chat requests so that each turn reuses KV cache blocks
@@ -54,7 +53,7 @@ Benchmarking:
it the benchmark sends OpenAI-schema-compliant payloads and every
turn lands as a cache MISS in this proxy.
Example:
Example:
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\
--model <MODEL> --served-model-name <NAME> \\
--url http://<proxy_host>:8000 \\
@@ -67,6 +66,7 @@ Benchmarking:
Dependencies:
pip install fastapi uvicorn httpx
"""
from __future__ import annotations
examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Push-mode disaggregated prefilling proxy demo.
"""Push-mode disaggregated prefilling proxy demo.
Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is
the same; the difference is in how P and D coordinate the KV transfer:
examples/disaggregated/example_connector/decode_example.py
+1 -1
View File
@@ -6,7 +6,7 @@ from vllm.config import KVTransferConfig
def read_prompts():
"""Read prompts from output.txt"""
"""Read prompts from output.txt."""
prompts = []
try:
with open("output.txt") as f:
examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use FlexKV with vLLM for prefix caching.
"""This example shows how to use FlexKV with vLLM for prefix caching.
FlexKV is a distributed KV Store and multi-level cache management system for
ultra-large-scale LLM inference.
examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py
+1 -1
View File
@@ -7,7 +7,7 @@ from vllm.config import KVTransferConfig
def read_prompts():
"""Read prompts from prefill_output.txt"""
"""Read prompts from prefill_output.txt."""
prompts = []
try:
with open("prefill_output.txt") as f:
examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py
+1 -2
View File
@@ -119,8 +119,7 @@ class LoadRecoveryExampleConnector(ExampleConnector):
def update_state_after_alloc(
self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
):
"""
Update KVConnector state after block allocation.
"""Update KVConnector state after block allocation.
If blocks were allocated, add to _requests_need_load,
such that we load the KVs in the next forward pass.
examples/disaggregated/lmcache/cpu_offload_lmcache.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of CPU offloading
"""This file demonstrates the example usage of CPU offloading
with LMCache in vLLM v1.
Note that `lmcache` is needed to run this example.
examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py
+3 -9
View File
@@ -14,9 +14,7 @@ from fastapi.responses import StreamingResponse
@asynccontextmanager
async def lifespan(app: FastAPI):
"""
Lifespan context manager to handle startup and shutdown events.
"""
"""Lifespan context manager to handle startup and shutdown events."""
# Startup: Initialize clients
prefiller_base_url = (
f"http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1"
@@ -106,9 +104,7 @@ app.state.decode_client = None
async def send_request_to_service(
client: httpx.AsyncClient, endpoint: str, req_data: dict
):
"""
Send a request to a service using a persistent client.
"""
"""Send a request to a service using a persistent client."""
req_data = req_data.copy()
req_data["max_tokens"] = 1
if "max_completion_tokens" in req_data:
@@ -128,9 +124,7 @@ async def send_request_to_service(
async def stream_service_response(
client: httpx.AsyncClient, endpoint: str, req_data: dict
):
"""
Asynchronously stream the response from a service using a persistent client.
"""
"""Asynchronously stream the response from a service using a persistent client."""
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
async with client.stream(
"POST", endpoint, json=req_data, headers=headers
examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of remote KV cache sharing
"""This file demonstrates the example usage of remote KV cache sharing
with LMCache.
We will launch 2 vllm instances, and launch an additional LMCache server.
KV cache is transferred in the following manner:
examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
+7 -11
View File
@@ -65,9 +65,7 @@ async def get_prefiller_info(prefill_clients: list, ready: asyncio.Event):
@asynccontextmanager
async def lifespan(app: FastAPI):
"""
Lifespan context manager to handle startup and shutdown events.
"""
"""Lifespan context manager to handle startup and shutdown events."""
# Startup: Initialize client pools for prefiller and decoder services
app.state.prefill_clients = []
app.state.decode_clients = []
@@ -180,10 +178,12 @@ def _parse_prefill_urls(prefill_list):
"""Parse prefill URLs from --prefill arguments.
Format: --prefill URL [BOOTSTRAP_PORT]
Example:
--prefill http://prefill1:8080 9000 # With bootstrap port
--prefill http://prefill2:8080 none # Explicitly no bootstrap port
--prefill http://prefill3:8080 # Defaults to no bootstrap port
"""
if not prefill_list:
return []
@@ -228,8 +228,7 @@ def _parse_decode_urls(decode_list):
def get_next_client(app, service_type: str):
"""
Get the next client in round-robin fashion.
"""Get the next client in round-robin fashion.
Args:
app: The FastAPI app instance
@@ -237,6 +236,7 @@ def get_next_client(app, service_type: str):
Returns:
The next client to use
"""
if service_type == "prefill":
return next(app.state.prefill_iterator)
@@ -250,9 +250,7 @@ def get_next_client(app, service_type: str):
async def send_request_to_service(
client_info: dict, dp_rank: int, endpoint: str, req_data: dict, request_id: str
):
"""
Send a request to a service using a client from the pool.
"""
"""Send a request to a service using a client from the pool."""
req_data = req_data.copy()
req_data["kv_transfer_params"] = {
"do_remote_decode": True,
@@ -288,9 +286,7 @@ async def stream_service_response(
req_data: dict,
request_id: str,
):
"""
Asynchronously stream response from a service using a client from the pool.
"""
"""Asynchronously stream response from a service using a client from the pool."""
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id,
examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstration script for Automatic Prefix Caching (APC) in vLLM.
"""Demonstration script for Automatic Prefix Caching (APC) in vLLM.
Automatic Prefix Caching (APC) allows the vLLM engine to reuse cached
KV (key-value) pairs from previous prompts if a new query shares the same
examples/features/batch_invariance/reproducibility_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates how to achieve reproducibility in vLLM.
"""Demonstrates how to achieve reproducibility in vLLM.
Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html
"""
examples/features/context_extension/context_extension_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This script demonstrates how to extend the context length
"""This script demonstrates how to extend the context length
of a Qwen model using the YARN method (rope_parameters)
and run a simple chat example.
examples/features/data_parallel/data_parallel_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Usage:
"""Usage:
Single node:
python examples/features/data_parallel/data_parallel_offline.py \
--model="ibm-research/PowerMoE-3b" \
examples/features/kv_events/kv_events_subscriber.py
+1 -1
View File
@@ -18,7 +18,7 @@ class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
"""Base class for all KV cache-related events"""
"""Base class for all KV cache-related events."""
class BlockStored(KVCacheEvent):
examples/features/logits_processor/custom.py
+1 -1
View File
@@ -48,7 +48,7 @@ from vllm.v1.sample.logits_processor.builtin import process_dict_updates
# Hypothetical custom logits processor
class DummyLogitsProcessor(LogitsProcessor):
"""Fake logit processor to support unit testing and examples"""
"""Fake logit processor to support unit testing and examples."""
@classmethod
def validate_params(cls, params: SamplingParams):
examples/features/logits_processor/custom_req.py
+6 -3
View File
@@ -55,10 +55,11 @@ logger = init_logger(__name__)
class DummyPerReqLogitsProcessor:
"""The request-level logits processor masks out all logits except the
token id identified by `target_token`"""
token id identified by `target_token`
"""
def __init__(self, target_token: int) -> None:
"""Specify `target_token`"""
"""Specify `target_token`."""
self.target_token = target_token
def __call__(
@@ -74,7 +75,8 @@ class DummyPerReqLogitsProcessor:
class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
"""Example of wrapping a fake request-level logit processor to create a
batch-level logits processor"""
batch-level logits processor
"""
@classmethod
def validate_params(cls, params: SamplingParams):
@@ -103,6 +105,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
Returns:
`Callable` request logits processor, or None
"""
target_token: Any | None = params.extra_args and params.extra_args.get(
"target_token"
examples/features/logits_processor/custom_req_init.py
+6 -3
View File
@@ -56,10 +56,11 @@ logger = init_logger(__name__)
class DummyPerReqLogitsProcessor:
"""The request-level logits processor masks out all logits except the
token id identified by `target_token`"""
token id identified by `target_token`
"""
def __init__(self, target_token: int) -> None:
"""Specify `target_token`"""
"""Specify `target_token`."""
self.target_token = target_token
def __call__(
@@ -75,7 +76,8 @@ class DummyPerReqLogitsProcessor:
class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
"""Example of overriding the wrapper class `__init__()` in order to utilize
info about the device type"""
info about the device type
"""
@classmethod
def validate_params(cls, params: SamplingParams):
@@ -111,6 +113,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
Returns:
`Callable` request logits processor, or None
"""
if (
not self.is_cuda
examples/features/lora/lora_with_quantization_offline.py
+1 -4
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use LoRA with different quantization techniques
"""This example shows how to use LoRA with different quantization techniques
for offline inference.
Requires HuggingFace credentials for access.
@@ -73,7 +72,6 @@ def initialize_engine(
model: str, quantization: str, lora_repo: str | None
) -> LLMEngine:
"""Initialize the LLMEngine."""
engine_args = EngineArgs(
model=model,
quantization=quantization,
@@ -86,7 +84,6 @@ def initialize_engine(
def main():
"""Main function that sets up and runs the prompt processing."""
test_configs = [
# QLoRA (https://arxiv.org/abs/2305.14314)
{
examples/features/lora/multilora_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use the multi-LoRA functionality
"""This example shows how to use the multi-LoRA functionality
for offline inference.
Requires HuggingFace credentials for access to Llama2.
examples/features/pause_resume/data_parallel_pause_resume.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Test pause/resume with Data Parallel (DP) via HTTP API.
"""Test pause/resume with Data Parallel (DP) via HTTP API.
This example demonstrates coordinated pause/resume across multiple DP ranks.
The pause synchronizes across all DP engines via all-reduce.
examples/features/prompt_embed/prompt_embed_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates how to generate prompt embeddings using
"""Demonstrates how to generate prompt embeddings using
Hugging Face Transformers and use them as input to vLLM
for both single and batch inference.
examples/features/reset_kv/reset_kv_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates preempt requests when using the `LLMEngine`
"""This file demonstrates preempt requests when using the `LLMEngine`
for processing prompts with various sampling parameters.
"""
examples/features/sharded_state/load_sharded_state_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Validates the loading of a model saved with the sharded_state format.
"""Validates the loading of a model saved with the sharded_state format.
This script demonstrates how to load a model that was previously saved
using save_sharded_state_offline.py and validates it by running inference.
Example usage:
examples/features/sharded_state/save_sharded_state_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Saves each worker's model state dict directly to a checkpoint, which enables a
"""Saves each worker's model state dict directly to a checkpoint, which enables a
fast load path for large tensor-parallel models where each worker only needs to
read its own shard rather than the entire checkpoint.
examples/features/speculative_decoding/mlpspeculator_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the usage of text generation with an LLM model,
"""This file demonstrates the usage of text generation with an LLM model,
comparing the performance with and without speculative decoding.
Note that this example is out of date and not supported in vLLM v1.
examples/features/structured_outputs/structured_outputs_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of structured outputs
"""This file demonstrates the example usage of structured outputs
in vLLM. It shows how to apply different constraints such as choice,
regex, json schema, and grammar to produce structured and formatted
results based on specific prompts.
examples/features/torchrun/torchrun_dp_example_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
experimental support for data-parallel inference with torchrun
"""experimental support for data-parallel inference with torchrun
Note the data load balancing and distribution is done out of the vllm engine,
no internal lb supported in external_launcher mode.
examples/features/torchrun/torchrun_example_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
experimental support for tensor-parallel inference with torchrun,
"""experimental support for tensor-parallel inference with torchrun,
see https://github.com/vllm-project/vllm/issues/11400 for
the motivation and use case for this example.
run the script with `torchrun --nproc-per-node=4 torchrun_example_offline.py`,
examples/generate/multimodal/audio_language_offline.py
+2 -4
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference
"""This example shows how to use vLLM for running offline inference
with the correct prompt format on audio language models.
For most models, the prompt format should follow corresponding examples
@@ -286,8 +285,7 @@ def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
# Phi-4-multimodal-instruct
def run_phi4mm(question: str, audio_count: int) -> ModelRequestData:
"""
Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
"""Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
show how to process audio inputs.
"""
model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
examples/generate/multimodal/encoder_decoder_multimodal_offline.py
+2 -4
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference with
"""This example shows how to use vLLM for running offline inference with
the explicit/implicit prompt format on enc-dec LMMs for text generation.
"""
@@ -56,8 +55,7 @@ def run_whisper():
def run_fireredasr2():
"""
FireRedASR2 – Automatic Speech Recognition model.
"""FireRedASR2 – Automatic Speech Recognition model.
This model uses a Conformer encoder + Qwen2 LLM decoder architecture
for speech-to-text transcription. Audio is passed via the implicit
examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
-2
View File
@@ -43,7 +43,6 @@ headers = {"User-Agent": "vLLM Example Client"}
def encode_base64_content_from_url(content_url: str) -> str:
"""Encode a content retrieved from a remote url to base64 format."""
with requests.get(content_url, headers=headers) as response:
response.raise_for_status()
result = base64.b64encode(response.content).decode("utf-8")
@@ -53,7 +52,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
def encode_base64_content_from_file(file_path: str) -> str:
"""Encode a local file content to base64 format."""
with open(file_path, "rb") as file:
file_content = file.read()
result = base64.b64encode(file_content).decode("utf-8")
examples/generate/multimodal/qwen2_5_omni/only_thinker.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference
"""This example shows how to use vLLM for running offline inference
with the correct prompt format on Qwen2.5-Omni (thinker only).
"""
examples/generate/multimodal/qwen3_omni/only_thinker.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference
"""This example shows how to use vLLM for running offline inference
with the correct prompt format on Qwen3-Omni (thinker only).
"""
examples/generate/multimodal/vision_language_multi_image_offline.py
+2 -5
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference with
"""This example shows how to use vLLM for running offline inference with
multi-image input on vision language models for text generation,
using the chat template defined by the model.
"""
@@ -913,11 +912,9 @@ def load_phi3v(question: str, image_urls: list[str]) -> ModelRequestData:
def load_phi4mm(question: str, image_urls: list[str]) -> ModelRequestData:
"""
Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
"""Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
show how to process multi images inputs.
"""
model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
# Since the vision-lora and speech-lora co-exist with the base model,
# we have to manually specify the path of the lora weights.
examples/generate/multimodal/vision_language_offline.py
+3 -6
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference with
"""This example shows how to use vLLM for running offline inference with
the correct prompt format on vision language models for text generation.
For most models, the prompt format should follow corresponding examples
@@ -1660,8 +1659,7 @@ def run_phi3v(questions: list[str], modality: str) -> ModelRequestData:
# Phi-4-multimodal-instruct
def run_phi4mm(questions: list[str], modality: str) -> ModelRequestData:
"""
Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
"""Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
show how to process image inputs.
"""
assert modality == "image"
@@ -2283,8 +2281,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
def get_multi_modal_input(args):
"""
return {
"""Return {
"data": image or video,
"question": question,
}
examples/generate/trace_replay_offline.py
+1 -2
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Trace-replay with vLLM offline inference.
"""Trace-replay with vLLM offline inference.
Trace-replay lets you supply a known sequence of decode token IDs alongside
the prompt. Instead of sampling from the model distribution, the engine
examples/pooling/embed/openai_embedding_long_text/client.py
+1 -2
View File
@@ -1,8 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example script demonstrating long text embedding with chunked processing in vLLM.
"""Example script demonstrating long text embedding with chunked processing in vLLM.
This example shows how to use vLLM's chunked processing feature to handle text
inputs that exceed the model's maximum token length. The feature automatically
examples/pooling/embed/vision_embedding_offline.py
+1 -2
View File
@@ -1,8 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ruff: noqa: E501
"""
This example shows how to use vLLM for running offline inference with
"""This example shows how to use vLLM for running offline inference with
the correct prompt format on vision language models for multimodal embedding.
For most models, the prompt format should follow corresponding examples

Some files were not shown because too many files have changed in this diff Show More